Code reference to web crawler video

web_crawler_gist.zig

const std = @import("std");
const Io = std.Io;

const stats_path = "crawl.stats";

const Job = struct { url: []const u8, depth: usize };

const State = struct {
    io: Io,
    gpa: std.mem.Allocator,
    arena: std.mem.Allocator,
    client: *std.http.Client,
    queue: Io.Queue(Job),
    max_pages: usize,
    max_depth: usize,

    started: std.atomic.Value(usize) = .init(0),
    pending: std.atomic.Value(usize) = .init(0),
    visited: std.atomic.Value(usize) = .init(0),
};

pub fn main(init: std.process.Init) !void {
    const gpa = init.gpa;
    const arena = init.arena.allocator();
    const args = try init.minimal.args.toSlice(arena);

    if (args.len < 2) {
        std.debug.print("usage: {s} <url> [max_pages] [workers] [max_depth]\n", .{args[0]});
        return;
    }

    const max_pages = if (args.len > 2) @max(try std.fmt.parseInt(usize, args[2], 10), 1) else 50;
    const workers = if (args.len > 3) @max(try std.fmt.parseInt(usize, args[3], 10), 1) else 8;
    const max_depth = if (args.len > 4) try std.fmt.parseInt(usize, args[4], 10) else 3;

    var client: std.http.Client = .{ .allocator = gpa, .io = init.io };
    defer client.deinit();

    const queue_buffer = try gpa.alloc(Job, max_pages);
    defer gpa.free(queue_buffer);

    var state: State = .{
        .io = init.io,
        .gpa = gpa,
        .arena = arena,
        .client = &client,
        .queue = .init(queue_buffer),
        .max_pages = max_pages,
        .max_depth = max_depth,
    };
    enqueue(&state, try arena.dupe(u8, args[1]), 0);

    var group: Io.Group = .init;
    defer group.cancel(init.io);

    for (0..@min(workers, max_pages)) |_| try group.concurrent(init.io, worker, .{&state});
    try group.await(init.io);
}

fn worker(s: *State) Io.Cancelable!void {
    while (true) {
        const job = s.queue.getOne(s.io) catch |err| switch (err) {
            error.Closed => return,
            error.Canceled => return error.Canceled,
        };
        defer done(s, job.url, job.depth);

        const html = fetch(s.gpa, s.client, job.url) catch |err| {
            std.log.warn("{s}: {t}", .{ job.url, err });
            continue;
        };
        defer s.gpa.free(html);

        _ = s.visited.fetchAdd(1, .monotonic);
        writeStats(s, "visited", job.url, job.depth, html.len) catch {};

        if (job.depth >= s.max_depth) continue;

        var pos: usize = 0;
        while (nextHref(html, &pos)) |href| {
            const url = resolve(s.arena, job.url, href) catch continue;
            enqueue(s, url, job.depth + 1);
        }
    }
}

fn enqueue(s: *State, url: []const u8, depth: usize) void {
    while (true) {
        const started = s.started.load(.monotonic);
        if (started >= s.max_pages) return;
        if (s.started.cmpxchgWeak(started, started + 1, .monotonic, .monotonic) == null) break;
    }

    _ = s.pending.fetchAdd(1, .monotonic);
    s.queue.putOneUncancelable(s.io, .{ .url = url, .depth = depth }) catch unreachable;
    writeStats(s, "found", url, depth, 0) catch {};
}

fn done(s: *State, url: []const u8, depth: usize) void {
    if (s.pending.fetchSub(1, .monotonic) == 1) {
        s.queue.close(s.io);
    }
    writeStats(s, "done", url, depth, 0) catch {};
}

fn writeStats(s: *State, event: []const u8, url: []const u8, depth: usize, bytes: usize) !void {
    var buf: [4096]u8 = undefined;
    const text = try std.fmt.bufPrint(&buf,
        \\pages found:   {d}/{d}
        \\pages visited: {d}
        \\in flight:     {d}
        \\max depth:     {d}
        \\last event:    {s}
        \\last depth:    {d}
        \\last bytes:    {d}
        \\last url:      {s}
        \\
    , .{
        s.started.load(.monotonic),
        s.max_pages,
        s.visited.load(.monotonic),
        s.pending.load(.monotonic),
        s.max_depth,
        event,
        depth,
        bytes,
        url,
    });

    try Io.Dir.cwd().writeFile(s.io, .{ .sub_path = stats_path, .data = text });
}

fn fetch(gpa: std.mem.Allocator, client: *std.http.Client, url: []const u8) ![]u8 {
    var body: Io.Writer.Allocating = .init(gpa);
    defer body.deinit();

    const result = try client.fetch(.{ .location = .{ .url = url }, .response_writer = &body.writer });
    std.log.info("{d} {s}", .{ @intFromEnum(result.status), url });
    if (result.status.class() != .success) return error.BadStatus;

    return body.toOwnedSlice();
}

fn resolve(gpa: std.mem.Allocator, base_url: []const u8, href: []const u8) ![]u8 {
    const clean = std.mem.trim(u8, href, &std.ascii.whitespace);
    if (clean.len == 0 or clean[0] == '#') return error.Skip;

    var buf: [4096]u8 = undefined;
    if (clean.len > buf.len) return error.Skip;
    @memcpy(buf[0..clean.len], clean);

    var rest: []u8 = &buf;
    const uri = try (try std.Uri.parse(base_url)).resolveInPlace(clean.len, &rest);

    const http =
        std.ascii.eqlIgnoreCase(uri.scheme, "http") or
        std.ascii.eqlIgnoreCase(uri.scheme, "https");
    if (!http) return error.Skip;

    return std.fmt.allocPrint(gpa, "{f}", .{&uri});
}

fn nextHref(html: []const u8, pos: *usize) ?[]const u8 {
    const hit = std.mem.cut(u8, html[pos.*..], "href=\"") orelse return null;
    const after = hit[1];
    const href, const rest = std.mem.cutScalar(u8, after, '"') orelse return null;
    pos.* = html.len - rest.len;
    return href;
}

test "href demo parser" {
    const html = "<a href=\"/a\"></a><a href=\"/b\"></a>";
    var pos: usize = 0;
    try std.testing.expectEqualStrings("/a", nextHref(html, &pos).?);
    try std.testing.expectEqualStrings("/b", nextHref(html, &pos).?);
}