Performing concurrent writes (to STDOUT) using a lock and a queue, in Zig 0.16.0

concurrent_writes.zig

const std = @import("std");

const LogMessage = struct {
    worker_id: usize,
    msg: []const u8,
};

const LogQueue = struct {
    mutex: std.Io.Mutex = .init,
    list: std.ArrayListUnmanaged(LogMessage) = .empty,

    pub fn push(self: *LogQueue, io: std.Io, allocator: std.mem.Allocator, item: LogMessage) !void {
        try self.mutex.lock(io);
        defer self.mutex.unlock(io);
        try self.list.append(allocator, item);
    }

    pub fn pop(self: *LogQueue, io: std.Io) ?LogMessage {
        self.mutex.lock(io) catch return null;
        defer self.mutex.unlock(io);
        if (self.list.items.len == 0) return null;
        return self.list.orderedRemove(0);
    }

    pub fn deinit(self: *LogQueue, allocator: std.mem.Allocator) void {
        for (self.list.items) |item| {
            allocator.free(item.msg);
        }
        self.list.deinit(allocator);
    }
};

const MutexWorkerCtx = struct {
    io: std.Io,
    stdout: std.Io.File,
    mutex: *std.Io.Mutex,
    worker_id: usize,
};

fn mutexWorkerTask(ctx: MutexWorkerCtx) void {

    ctx.io.sleep(.fromMilliseconds(@intCast(ctx.worker_id * 30)), .real) catch return;

    var buf: [128]u8 = undefined;
    const line = std.fmt.bufPrint(&buf, "[Mutex Worker {}] Wrote directly via lock!\n", .{ctx.worker_id}) catch return;

    ctx.mutex.lock(ctx.io) catch return;
    defer ctx.mutex.unlock(ctx.io);

    ctx.stdout.writeStreamingAll(ctx.io, line) catch return;
}

const QueueWorkerCtx = struct {
    io: std.Io,
    allocator: std.mem.Allocator,
    queue: *LogQueue,
    worker_id: usize,
};

fn queueWorkerTask(ctx: QueueWorkerCtx) void {

    ctx.io.sleep(.fromMilliseconds(@intCast(ctx.worker_id * 30)), .real) catch return;

    const msg = std.fmt.allocPrint(ctx.allocator, "Processed data unit #{}", .{ctx.worker_id * 10}) catch return;

    ctx.queue.push(ctx.io, ctx.allocator, .{
        .worker_id = ctx.worker_id,
        .msg = msg,
    }) catch {
        ctx.allocator.free(msg);
    };
}

pub fn main(init: std.process.Init) !void {
    const allocator = std.heap.page_allocator;
    const io = init.io;
    const stdout = std.Io.File.stdout();

    // ==========================================
    // Method one: using a lock
    // ==========================================
    try stdout.writeStreamingAll(io, "--- Starting (Methd One) ---\n");

    var log_mutex: std.Io.Mutex = .init;
    var mutex_threads: std.ArrayList(std.Thread) = .empty;
    defer mutex_threads.deinit(allocator);

    var i: usize = 0;
    while (i < 3) : (i += 1) {
        const ctx = MutexWorkerCtx{
            .io = io,
            .stdout = stdout,
            .mutex = &log_mutex,
            .worker_id = i + 1,
        };
        const thread = try std.Thread.spawn(.{}, mutexWorkerTask, .{ctx});
        try mutex_threads.append(allocator, thread);
    }

    for (mutex_threads.items) |thread| thread.join();

    // ==========================================
    // Method two: using a queue
    // ==========================================
    try stdout.writeStreamingAll(io, "\n--- Starting (Method Two) ---\n");

    var log_queue = LogQueue{};
    defer log_queue.deinit(allocator);

    var queue_threads: std.ArrayList(std.Thread) = .empty;
    defer queue_threads.deinit(allocator);

    i = 0;
    while (i < 3) : (i += 1) {
        const ctx = QueueWorkerCtx{
            .io = io,
            .allocator = allocator,
            .queue = &log_queue,
            .worker_id = i + 1,
        };
        const thread = try std.Thread.spawn(.{}, queueWorkerTask, .{ctx});
        try queue_threads.append(allocator, thread);
    }

    for (queue_threads.items) |thread| thread.join();

    var buf: [256]u8 = undefined;
    while (log_queue.pop(io)) |item| {
        defer allocator.free(item.msg);
        const line = try std.fmt.bufPrint(&buf, "[Queue Consumer] Worker {} sent message: {s}\n", .{ item.worker_id, item.msg });
        try stdout.writeStreamingAll(io, line);
    }
}