mpmc_bus.zig

const std = @import("std");
const assert = std.debug.assert;

const stdx = @import("stdx");
const BufferedChannel = stdx.BufferedChannel;
const UnbufferedChannel = stdx.UnbufferedChannel;
const CancellationToken = stdx.CancellationToken;
const RingBuffer = stdx.RingBuffer;
const log = std.log.scoped(.MPSCBusExample);

const Dog = struct {
    name: []const u8,
    age: u8,
};

const sardine = Dog{
    .name = "sardine",
    .age = 4,
};

// This is the type that will be processed
const VALUE_TYPE: type = *const Dog;
const BUS_QUEUE_SIZE = 100_000;
const PRODUCER_QUEUE_SIZE = 10_000;
const CONSUMER_QUEUE_SIZE = 1_000;
const ITERATIONS = 10_000;
const CONSUMER_COUNT = 100;
const PRODUCER_COUNT = 10;

pub fn doProduce(
    io: std.Io,
    producers: *std.array_list.Managed(*Producer(VALUE_TYPE)),
    iterations: usize,
    value: VALUE_TYPE,
    ready_channel: *UnbufferedChannel(bool),
) void {
    ready_channel.send(true);
    for (0..iterations) |_| {
        var i: usize = 0;
        while (i < producers.items.len) {
            const producer = producers.items[i];

            // try to produce an item
            producer.produce(value) catch {
                // this producer is too fast
                std.Io.sleep(io, .fromMilliseconds(1), .awake) catch unreachable;

                continue;
            };

            i += 1;
        }
    }
}

pub fn Bus(comptime T: type) type {
    return struct {
        const Self = @This();

        allocator: std.mem.Allocator,
        io: std.Io,
        mutex: std.Io.Mutex,
        queue: *RingBuffer(T),
        consumers: *std.array_list.Managed(*Consumer(T)),
        producers: *std.array_list.Managed(*Producer(T)),
        close_channel: UnbufferedChannel(bool),
        last_producer_index: usize,

        pub fn init(allocator: std.mem.Allocator, io: std.Io) !Self {
            const queue = try allocator.create(RingBuffer(T));
            errdefer allocator.destroy(queue);

            queue.* = try RingBuffer(T).initCapacity(allocator, BUS_QUEUE_SIZE);
            errdefer queue.deinit(allocator);

            const consumers = try allocator.create(std.array_list.Managed(*Consumer(T)));
            errdefer allocator.destroy(consumers);

            consumers.* = std.array_list.Managed(*Consumer(T)).init(allocator);
            errdefer consumers.deinit();

            const producers = try allocator.create(std.array_list.Managed(*Producer(T)));
            errdefer allocator.destroy(producers);

            producers.* = std.array_list.Managed(*Producer(T)).init(allocator);
            errdefer producers.deinit();

            return Self{
                .allocator = allocator,
                .mutex = .init,
                .queue = queue,
                .consumers = consumers,
                .producers = producers,
                .close_channel = UnbufferedChannel(bool).new(io),
                .last_producer_index = 0,
                .io = io,
            };
        }

        pub fn deinit(self: *Self) void {
            self.queue.deinit(self.allocator);
            self.consumers.deinit();
            self.producers.deinit();

            self.allocator.destroy(self.queue);
            self.allocator.destroy(self.consumers);
            self.allocator.destroy(self.producers);
        }

        pub fn tick(self: *Self) !void {
            if (self.producers.items.len > 0) {
                self.mutex.lockUncancelable(self.io);
                defer self.mutex.unlock(self.io);

                var processed: usize = 0;
                const producer_count = self.producers.items.len;

                while (processed < producer_count) : (processed += 1) {
                    const producer_index = (self.last_producer_index + processed) % producer_count;
                    const producer = self.producers.items[producer_index];

                    if (self.queue.available() == 0) {
                        log.debug("bus queue full producer index: {}", .{producer_index});

                        // next tick should resume from the next producer
                        self.last_producer_index = producer_index;
                        break;
                    }

                    producer.mutex.lockUncancelable(self.io);
                    defer producer.mutex.unlock(self.io);

                    _ = self.queue.concatenateAvailable(producer.queue);
                }

                // If we completed the loop, set last index to the next producer
                if (processed == producer_count and self.queue.available() > 0) {
                    self.last_producer_index = (self.last_producer_index + 1) % producer_count;
                }
            }

            if (self.consumers.items.len > 0) {
                self.mutex.lockUncancelable(self.io);
                defer self.mutex.unlock(self.io);

                var max_available = self.queue.count;
                for (self.consumers.items) |consumer| {
                    const consumer_available = consumer.queue.available();
                    if (consumer_available < max_available) {
                        max_available = consumer_available;
                    }
                }
            }

            // if there are no items in the queue, then there is nothing to do
            if (self.queue.count == 0) return;

            // if there are no consumers of the items on the bus, then there is no work to be done
            if (self.consumers.items.len == 0) return;

            // FIX: there should be a consumer mutex to ensure that the number of consumers remains constant throughout this tick
            const consumer_queues = try self.allocator.alloc(*RingBuffer(T), self.consumers.items.len);
            defer self.allocator.free(consumer_queues);

            self.mutex.lockUncancelable(self.io);
            defer self.mutex.unlock(self.io);

            for (self.consumers.items, 0..self.consumers.items.len) |consumer, i| {
                consumer.mutex.lockUncancelable(self.io);
                consumer_queues[i] = consumer.queue;
            }
            defer {
                for (self.consumers.items) |consumer| {
                    consumer.mutex.unlock(self.io);
                }
            }

            _ = self.queue.copyMaxToOthers(consumer_queues);
        }

        pub fn run(self: *Self, ready: *UnbufferedChannel(bool)) void {
            ready.send(true);
            while (true) {
                // check if we have received a signale to close the topic
                const signal = self.close_channel.tryReceive(.fromMilliseconds(0)) catch false;
                if (signal) {
                    return;
                }

                self.tick() catch unreachable;
                std.Io.sleep(self.io, .fromMilliseconds(1), .awake) catch unreachable;
            }
        }

        pub fn close(self: *Self) void {
            self.close_channel.send(true);
        }
    };
}

pub fn Consumer(comptime T: type) type {
    return struct {
        const Self = @This();

        allocator: std.mem.Allocator,
        close_channel: UnbufferedChannel(bool),
        id: usize,
        mutex: std.Io.Mutex,
        consumed_count: u128,
        queue: *RingBuffer(T),
        bus: *Bus(T),
        io: std.Io,

        pub fn init(allocator: std.mem.Allocator, io: std.Io, id: usize, bus: *Bus(T)) !Self {
            const queue = try allocator.create(RingBuffer(T));
            errdefer allocator.destroy(queue);

            queue.* = try RingBuffer(T).initCapacity(allocator, CONSUMER_QUEUE_SIZE);
            errdefer queue.deinit(allocator);

            return Self{
                .allocator = allocator,
                .close_channel = UnbufferedChannel(bool).new(io),
                .id = id,
                .mutex = .init,
                .consumed_count = 0,
                .queue = queue,
                .bus = bus,
                .io = io,
            };
        }

        pub fn deinit(self: *Self) void {
            self.queue.deinit(self.allocator);

            self.allocator.destroy(self.queue);
        }

        pub fn tick(self: *Self) !void {
            if (self.queue.count == 0) return;

            self.mutex.lockUncancelable(self.io);
            defer self.mutex.unlock(self.io);

            // FIX: this is just some BS where the consumer is dropping the items and not
            // actually doing any work. This is an example so don't look to deeply into it.
            self.consumed_count += self.queue.count;
            self.queue.reset();
        }

        pub fn run(self: *Self, ready: *UnbufferedChannel(bool)) void {
            ready.send(true);
            while (true) {
                // check if we have received a signale to close the topic
                const signal = self.close_channel.tryReceive(.fromMilliseconds(0)) catch false;
                if (signal) {
                    return;
                }

                self.tick() catch unreachable;
                std.Io.sleep(self.io, .fromMilliseconds(1), .awake) catch unreachable;
            }
        }

        pub fn close(self: *Self) void {
            self.close_channel.send(true);
        }
    };
}

pub fn Producer(comptime T: type) type {
    return struct {
        const Self = @This();

        allocator: std.mem.Allocator,
        close_channel: UnbufferedChannel(bool),
        id: usize,
        mutex: std.Io.Mutex,
        produced_count: u128,
        queue: *RingBuffer(T),
        io: std.Io,

        pub fn init(allocator: std.mem.Allocator, io: std.Io, id: usize) !Self {
            const queue = try allocator.create(RingBuffer(T));
            errdefer allocator.destroy(queue);

            queue.* = try RingBuffer(T).initCapacity(allocator, PRODUCER_QUEUE_SIZE);
            errdefer queue.deinit(allocator);

            return Self{
                .allocator = allocator,
                .close_channel = UnbufferedChannel(bool).new(io),
                .id = id,
                .mutex = .init,
                .produced_count = 0,
                .queue = queue,
                .io = io,
            };
        }

        pub fn deinit(self: *Self) void {
            self.queue.deinit(self.allocator);
            self.allocator.destroy(self.queue);
        }

        pub fn produce(self: *Self, value: T) !void {
            self.mutex.lockUncancelable(self.io);
            defer self.mutex.unlock(self.io);

            try self.queue.enqueue(self.allocator, value);

            self.produced_count += 1;
        }
    };
}

pub fn main(init: std.process.Init) !void {
    const io = init.io;
    var gpa = std.heap.DebugAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    var bus = try Bus(VALUE_TYPE).init(allocator, io);
    defer bus.deinit();

    var consumers = std.array_list.Managed(*Consumer(VALUE_TYPE)).init(allocator);
    defer consumers.deinit();

    var producers = std.array_list.Managed(*Producer(VALUE_TYPE)).init(allocator);
    defer producers.deinit();

    var bus_ready_channel = UnbufferedChannel(bool).new(io);

    const th = try std.Thread.spawn(
        .{},
        Bus(VALUE_TYPE).run,
        .{ &bus, &bus_ready_channel },
    );
    th.detach();

    _ = bus_ready_channel.receive();
    defer bus.close();

    try std.Io.sleep(io, .fromMilliseconds(500), .awake);

    // spawn all the consumers
    for (0..CONSUMER_COUNT) |i| {
        const consumer = try allocator.create(Consumer(VALUE_TYPE));
        errdefer allocator.destroy(consumer);

        consumer.* = try Consumer(VALUE_TYPE).init(allocator, io, i, &bus);
        errdefer consumer.deinit();

        try consumers.append(consumer);

        var ready_channel = UnbufferedChannel(bool).new(io);

        const consumer_thread = try std.Thread.spawn(
            .{},
            Consumer(VALUE_TYPE).run,
            .{ consumer, &ready_channel },
        );
        consumer_thread.detach();

        _ = ready_channel.receive();
        errdefer consumer.close();

        bus.mutex.lockUncancelable(io);
        defer bus.mutex.unlock(io);

        try bus.consumers.append(consumer);
    }

    // spawn all the producers
    for (0..PRODUCER_COUNT) |i| {
        const producer = try allocator.create(Producer(VALUE_TYPE));
        errdefer allocator.destroy(producer);

        producer.* = try Producer(VALUE_TYPE).init(allocator, io, i);
        errdefer producer.deinit();

        try producers.append(producer);

        bus.mutex.lockUncancelable(io);
        defer bus.mutex.unlock(io);

        try bus.producers.append(producer);
    }

    // spawn all the consumers
    assert(bus.producers.items.len == PRODUCER_COUNT);
    assert(bus.consumers.items.len == CONSUMER_COUNT);

    var do_produce_ready_chan = UnbufferedChannel(bool).new(io);
    const do_produce_thread = try std.Thread.spawn(.{}, doProduce, .{
        io,
        &producers,
        ITERATIONS,
        &sardine,
        &do_produce_ready_chan,
    });
    do_produce_thread.detach();

    _ = do_produce_ready_chan.receive();

    // everything is setup now
    const start = std.Io.Timestamp.now(io, .awake);

    var total_items_produced: u128 = 0;
    while (total_items_produced != ITERATIONS * PRODUCER_COUNT) {
        std.Io.sleep(io, .fromMilliseconds(1), .awake) catch unreachable;
        total_items_produced = 0;
        for (producers.items) |producer| {
            total_items_produced += producer.produced_count;
        }
    }

    var total_items_consumed: u128 = 0;
    while (total_items_consumed != ITERATIONS * PRODUCER_COUNT * CONSUMER_COUNT) {
        std.Io.sleep(io, .fromMilliseconds(1), .awake) catch unreachable;
        total_items_consumed = 0;
        for (consumers.items) |consumer| {
            total_items_consumed += consumer.consumed_count;
        }
    }

    const end = std.Io.Timestamp.now(io, .awake);
    log.err("took {d}ms, total iters {}, total_items_produced {}, total_items_consumed {}", .{
        @divTrunc(end.nanoseconds - start.nanoseconds, std.time.ns_per_ms),
        ITERATIONS,
        total_items_produced,
        total_items_consumed,
    });

    log.err("producer_count {}, consumer count {}", .{
        PRODUCER_COUNT,
        CONSUMER_COUNT,
    });

    // Clean up all of the producers
    for (producers.items) |producer| {
        // producer.close();
        producer.deinit();
        allocator.destroy(producer);
    }

    // Clean up all of the consumers
    for (consumers.items) |consumer| {
        // remove this consumer from the bus
        bus.mutex.lockUncancelable(io);
        defer bus.mutex.unlock(io);

        for (bus.consumers.items, 0..bus.consumers.items.len) |bus_consumer, i| {
            if (consumer == bus_consumer) {
                _ = bus.consumers.swapRemove(i);
                break;
            }
        }

        consumer.close();
        consumer.deinit();
        allocator.destroy(consumer);
    }
}