castle/packages/web/src/Server.zig

const std = @import("std");
const Server = @This();

const Address = @import("Address.zig").Address;
const Connection = @import("Connection.zig");
const FileDescriptor = @import("FileDescriptor.zig").FileDescriptor;
const http = @import("http.zig");
const openssl = @import("openssl.zig");
const Request = @import("Request.zig");
const RequestHandler = @import("RequestHandler.zig");
const Worker = @import("Worker.zig");

const log = std.log.scoped(.Server);
const linux = std.os.linux;
const errno = linux.errno;

fd: FileDescriptor,
address: Address,
ssl_ctx: ?*openssl.SslContext,
workers: []Worker,
threads: []std.Thread,
request_handler: RequestHandler,
read_timeout_us: u64,

connection_queue: std.DoublyLinkedList,
// NOTE Connection pool has no need for being doubly-linked, but the queue has
// (as it's FIFO) and we want a single intrusive Node to be able to participate
// in both lists. This is possible because a connection will never belong to
// both lists at the same time.
connection_pool: std.DoublyLinkedList,
connection_buffer: []Connection,

io: std.Io,
mutex: std.Io.Mutex,
cond_connection_queued: std.Io.Condition,
cond_connection_freed: std.Io.Condition,

/// 4 kiB
const page_size = 4 * 1024;
/// 2 MiB
const huge_page_size = 2 * 1024 * 1024;

pub const Options = struct {
    request_handler: RequestHandler,
    address: Address = .initIp4(.{ 127, 0, 0, 1 }, 8000),
    /// If not `null`, the server will use TLS with the provided OpenSSL
    /// context.
    ssl_ctx: ?*openssl.SslContext = null,
    max_connections: u32 = 128,
    /// The number of worker threads. If set to `0`, the number of worker
    /// threads will be equal to the number of logical CPU cores.
    worker_count: u32 = 0,
    /// The maximum number of header fields the `Request` object will be able to
    /// store. An HTTP request will be rejected if it has more header fields
    /// than the capacity.
    max_header_fields: u32 = 256,
    /// The number of 2 MiB pages reserved for a single read buffer. Each worker
    /// has its own read buffer. An HTTP request (headers and content combined)
    /// will be rejected if it is larger than the read buffer.
    read_buffer_huge_pages: u32 = 1,
    /// The number of 4 kiB pages reserved for a single header write buffer.
    /// Each worker has its own header write buffer. The HTTP status line, all
    /// header fields and the CRLF terminator must all fit in the header write
    /// buffer.
    header_write_buffer_pages: u32 = 1,
    /// The number of 2 MiB pages reserved for a single body write buffer. Each
    /// worker has its own body write buffer. The HTTP response body must fit
    /// entirely within the body write buffer. This restriction only applies to
    /// bodies generated with the body writer and not to bodies sent with
    /// `sendfile`.
    body_write_buffer_huge_pages: u32 = 1,
    /// How much time should a worker wait on an idle connection before closing
    /// it. Specifically, how much time can a `read` syscall block for, before
    /// the connection is forcefully closed.
    read_timeout_us: u64 = 1 * std.time.us_per_s,
};

pub fn init(allocator: std.mem.Allocator, io: std.Io, options: Options) !Server {
    const worker_count = if (options.worker_count > 0) options.worker_count else try std.Thread.getCpuCount();

    // Create socket fd

    const fd: FileDescriptor = try .socket(
        options.address.any.family,
        linux.SOCK.STREAM | linux.SOCK.CLOEXEC,
        if (options.address.any.family == linux.AF.UNIX) 0 else linux.IPPROTO.TCP,
    );
    errdefer fd.close();

    const opt = std.mem.toBytes(@as(c_int, 1));
    try fd.setsockopt(linux.SOL.SOCKET, linux.SO.REUSEADDR, &opt);
    try fd.setsockopt(linux.SOL.SOCKET, linux.SO.REUSEPORT, &opt);

    var socklen = options.address.getOsSockLen();
    try fd.bind(&options.address.any, socklen);
    try fd.listen(options.max_connections);

    var listen_address = options.address;
    try fd.getsockname(&listen_address.any, &socklen);

    // Allocate arrays

    const workers = try allocator.alloc(Worker, worker_count);
    errdefer allocator.free(workers);

    const connection_buffer = try allocator.alloc(Connection, options.max_connections);
    errdefer allocator.free(connection_buffer);

    const threads = try allocator.alloc(std.Thread, worker_count);
    errdefer allocator.free(threads);

    // Allocate and remap read buffers

    const single_read_buffer_size = @as(usize, options.read_buffer_huge_pages) * huge_page_size;
    const all_read_buffers_size = worker_count * single_read_buffer_size;

    const double_single_read_buffer_size = 2 * single_read_buffer_size;
    const double_all_read_buffers_size = 2 * all_read_buffers_size;

    const read_buffer_fd: FileDescriptor = try .memfd_create("read_buffer", 0);
    defer read_buffer_fd.close();

    try read_buffer_fd.ftruncate(@intCast(all_read_buffers_size));

    const read_buffer_ptr = try errOrPtr(linux.mmap(
        null,
        double_all_read_buffers_size,
        .{},
        linux.MAP{ .TYPE = .PRIVATE, .ANONYMOUS = true },
        -1,
        0,
    ));
    errdefer _ = linux.munmap(read_buffer_ptr, double_all_read_buffers_size);
    _ = linux.madvise(read_buffer_ptr, double_all_read_buffers_size, linux.MADV.HUGEPAGE);

    for (0..worker_count) |i| {
        const offset = i * single_read_buffer_size;
        const double_offset = i * double_single_read_buffer_size;

        try err(linux.mmap(
            read_buffer_ptr + double_offset,
            single_read_buffer_size,
            .{ .READ = true, .WRITE = true },
            .{ .TYPE = .SHARED, .FIXED = true },
            @intFromEnum(read_buffer_fd),
            @intCast(offset),
        ));

        try err(linux.mmap(
            read_buffer_ptr + double_offset + single_read_buffer_size,
            single_read_buffer_size,
            .{ .READ = true, .WRITE = true },
            .{ .TYPE = .SHARED, .FIXED = true },
            @intFromEnum(read_buffer_fd),
            @intCast(offset),
        ));
    }

    // Allocate header write buffer

    const single_header_write_buffer_size = @as(usize, options.header_write_buffer_pages) * page_size;
    const all_header_write_buffers_size = worker_count * single_header_write_buffer_size;

    const header_write_buffer_ptr = try errOrPtr(linux.mmap(
        null,
        all_header_write_buffers_size,
        .{ .READ = true, .WRITE = true },
        .{ .TYPE = .PRIVATE, .ANONYMOUS = true },
        -1,
        0,
    ));
    errdefer _ = linux.munmap(header_write_buffer_ptr, all_header_write_buffers_size);

    // Allocate body write buffer

    const single_body_write_buffer_size = @as(usize, options.body_write_buffer_huge_pages) * huge_page_size;
    const all_body_write_buffers_size = worker_count * single_body_write_buffer_size;

    const body_write_buffer_ptr = try errOrPtr(linux.mmap(
        null,
        all_body_write_buffers_size,
        .{ .READ = true, .WRITE = true },
        .{ .TYPE = .PRIVATE, .ANONYMOUS = true },
        -1,
        0,
    ));
    errdefer _ = linux.munmap(body_write_buffer_ptr, all_body_write_buffers_size);
    _ = linux.madvise(body_write_buffer_ptr, all_body_write_buffers_size, linux.MADV.HUGEPAGE);

    // Initialize workers

    var workers_initialized: usize = 0;
    errdefer {
        for (workers[0..workers_initialized]) |*worker| {
            worker.deinit(allocator);
        }
    }

    for (workers, 0..) |*worker, i| {
        const read_offset = i * double_single_read_buffer_size;
        const header_write_offset = i * single_header_write_buffer_size;
        const body_write_offset = i * single_body_write_buffer_size;
        worker.* = try Worker.init(allocator, .{
            .worker_id = i,

            .max_header_fields = options.max_header_fields,

            .read_buffer_ptr = read_buffer_ptr + read_offset,
            .read_buffer_size = single_read_buffer_size,

            .header_write_buffer = (header_write_buffer_ptr + header_write_offset)[0..single_header_write_buffer_size],
            .body_write_buffer = (body_write_buffer_ptr + body_write_offset)[0..single_body_write_buffer_size],
        });
        workers_initialized += 1;
    }

    // Fill connection pool

    var connection_pool: std.DoublyLinkedList = .{};
    for (connection_buffer) |*c| {
        connection_pool.prepend(&c.node);
    }

    return .{
        .fd = fd,
        .address = listen_address,
        .ssl_ctx = options.ssl_ctx,
        .workers = workers,
        .threads = threads,
        .request_handler = options.request_handler,
        .read_timeout_us = options.read_timeout_us,

        .connection_queue = .{},
        .connection_pool = connection_pool,
        .connection_buffer = connection_buffer,

        .io = io,
        .mutex = .init,
        .cond_connection_queued = .init,
        .cond_connection_freed = .init,
    };
}

pub fn deinit(self: *Server, allocator: std.mem.Allocator) void {
    log.debug("Deinitializing Server.", .{});
    const worker_count = self.workers.len;

    const single_read_buffers_size = self.workers[0].read_buffer_size;
    const all_read_buffers_size = worker_count * single_read_buffers_size;
    const double_all_read_buffers_size = 2 * all_read_buffers_size;

    const single_header_write_buffer_size = self.workers[0].header_write_buffer.len;
    const all_header_write_buffers_size = worker_count * single_header_write_buffer_size;

    const single_body_write_buffer_size = self.workers[0].body_write_buffer.len;
    const all_body_write_buffers_size = worker_count * single_body_write_buffer_size;

    const read_buffer_ptr = self.workers[0].read_buffer_ptr;
    const header_write_buffer_ptr = self.workers[0].header_write_buffer.ptr;
    const body_write_buffer_ptr = self.workers[0].body_write_buffer.ptr;

    for (self.workers) |*worker| {
        worker.deinit(allocator);
    }

    _ = linux.munmap(body_write_buffer_ptr, all_body_write_buffers_size);
    _ = linux.munmap(header_write_buffer_ptr, all_header_write_buffers_size);
    _ = linux.munmap(read_buffer_ptr, double_all_read_buffers_size);

    allocator.free(self.threads);
    allocator.free(self.connection_buffer);
    allocator.free(self.workers);

    self.fd.close();

    self.* = undefined;
}

/// This method block until the server is stopped, which is achieved by storing
/// `false` into `running`. You should use another thread or interruption
/// handler to be able to stop the server.
pub fn listen(self: *Server, running: *const std.atomic.Value(bool)) !void {
    var worker_running: std.atomic.Value(bool) = .init(running.load(.acquire));

    var spawned: usize = 0;
    defer {
        log.debug("Storing `false` into worker_running.", .{});
        worker_running.store(false, .release);
        log.debug("Broadcasting connection queued condition variable.", .{});
        self.cond_connection_queued.broadcast(self.io);
        for (self.threads[0..spawned], 0..) |*thread, i| {
            log.debug("Joining the thread of worker #{d}.", .{i});
            thread.join();
        }
    }

    for (self.workers, 0..) |*worker, i| {
        log.debug("Spawning thread for worker #{d}.", .{i});
        self.threads[i] = try std.Thread.spawn(.{}, Worker.worker, .{ worker, self, &worker_running });
        spawned += 1;
    }

    while (running.load(.acquire)) {
        var address: Address = undefined;
        var address_size: u32 = @sizeOf(Address);

        log.debug("Accepting connection.", .{});
        const fd = self.fd.accept(&address.any, &address_size) catch |e| {
            log.err("Error while accepting connection: {}", .{e});
            continue;
        };
        log.debug("Accepted connection from {f}", .{address});

        const timeout: linux.timeval = .{
            .sec = @intCast(self.read_timeout_us / std.time.us_per_s),
            .usec = @intCast(self.read_timeout_us % std.time.us_per_s),
        };
        try fd.setsockopt(linux.SOL.SOCKET, linux.SO.RCVTIMEO, std.mem.asBytes(&timeout));

        const ssl: ?*openssl.Ssl = self.maybeInitSsl(fd) catch |e| {
            log.err("Error while estabilishing SSL connection: {}", .{e});
            fd.close();
            continue;
        };

        {
            log.debug("Acquiring mutex.", .{});
            self.mutex.lockUncancelable(self.io);
            log.debug("Acquired mutex.", .{});
            defer {
                log.debug("Unlocking mutex.", .{});
                self.mutex.unlock(self.io);
            }

            while (true) {
                if (self.connection_pool.pop()) |node| {
                    const connection: *Connection = @fieldParentPtr("node", node);
                    connection.reinit(address, fd, ssl);
                    log.debug("Adding connection to {f} to the connection queue.", .{connection.address});
                    self.connection_queue.prepend(node);
                    break;
                }

                log.debug("Waiting on connection freed condition variable.", .{});
                self.cond_connection_freed.waitUncancelable(self.io, &self.mutex);
                log.debug("Woken up on connection freed condition variable.", .{});
            }
        }

        log.debug("Signaling connection queued condition variable.", .{});
        self.cond_connection_queued.signal(self.io);
    } else {
        log.debug("Loaded `false` from running, the accept loop exited.", .{});
    }
}

fn maybeInitSsl(self: *const Server, fd: FileDescriptor) !?*openssl.Ssl {
    if (self.ssl_ctx) |ssl_ctx| {
        const ssl = try openssl.Ssl.new(ssl_ctx);
        try ssl.setFd(fd);
        log.debug("Accepting SSL layer.", .{});
        try ssl.accept();
        log.debug("Accepted SSL layer.", .{});
        const opt = std.mem.toBytes(@as(c_int, 1));
        if (fd.setsockopt(linux.SOL.TLS, 3, &opt)) { // TLS_TX_ZEROCOPY_RO
            log.debug("Enabled zero-copy on sendfile.", .{});
        } else |_| {}
        return ssl;
    } else {
        return null;
    }
}

fn err(rc: usize) !void {
    const e = errno(rc);
    return if (e != .SUCCESS) error.SystemError else {};
}

fn errOrPtr(rc: usize) ![*]u8 {
    const e = errno(rc);
    return if (e != .SUCCESS) error.SystemError else @ptrFromInt(rc);
}