Lock-free SPSC Triple Buffer Mailbox

Hello everyone, I’d like to share this lock-free single-producer single-consumer mailbox implementation I wrote.

My use case is that I need to pass messages from a high real-time priority thread to a low non-real-time priority thread without any blocking of the high prio thread. I didn’t want a mutex because that would depend on the OS / std lib correctly implementing priority inheritance for mutexes, which std.Io.Threaded does not currently support. And in general, if I can rely on fewer features of the OS I try to.

Some cool things:

  • generic, supports any type to be buffered
  • two phase put and get to support zero-copy (producer / consumer have direct access to their buffers)
  • internally uses enum indexing into the backing triple buffer
  • atomic operations on a packed struct, much better than bit shifting stuff I saw a lot in cpp implementations

Areas for improvement:

  • enforce minimum cache line alignment on some of the fields
  • fuzz test
  • refactor API to provide non-optional on consumerPtr :person_shrugging:
  • relax the atomic ordering once I actually understand atomic ordering
const std = @import("std");
const assert = std.debug.assert;

pub fn SPSCTripleBuffer(T: type) type {
    return struct {
        bufs: [3]T = undefined,
        shared_state: std.atomic.Value(State) = .init(.{}),
        producer_slot: Slot = .back, // only touched by producer
        consumer_slot: Slot = .front, // only touched by consumer
        consumer_done: bool = true,

        // u8 since must be extern compatible
        pub const State = packed struct(u8) {
            has_data: bool = false,
            spare_slot: Slot = .middle,
            _reserved: u5 = 0,
        };

        pub const Slot = enum(u2) {
            back = 0,
            middle = 1,
            front = 2,

            pub fn toIndex(self: Slot) usize {
                return @as(usize, @as(u2, @intFromEnum(self)));
            }
        };

        /// Returns the same ptr for all calls before producerCommit.
        /// Invalidated on producerCommit.
        pub fn producerPtr(self: *@This()) *T {
            return &self.bufs[self.producer_slot.toIndex()];
        }

        /// Only call if you have previously finished writing data using producerPtr,
        /// otherwise undefined data may be sent to consumer.
        pub fn producerCommit(self: *@This()) void {
            const old_state = self.shared_state.swap(.{ .has_data = true, .spare_slot = self.producer_slot }, .seq_cst);
            assert(self.producer_slot != old_state.spare_slot);
            self.producer_slot = old_state.spare_slot;
        }

        /// producer only, non-blocking lock-free
        pub fn put(self: *@This(), value: T) void {
            self.producerPtr().* = value;
            self.producerCommit();
        }

        /// Returns null when nothing to consume.
        /// If non-null, returns same value for all calls before consumerCommit.
        /// Invalidated on consumerCommit.
        pub fn consumerPtr(self: *@This()) ?*T {
            if (!self.consumer_done) return &self.bufs[self.consumer_slot.toIndex()];
            assert(self.consumer_done);
            if (!self.shared_state.load(.seq_cst).has_data) return null;
            const old_state = self.shared_state.swap(.{ .has_data = false, .spare_slot = self.consumer_slot }, .seq_cst);
            assert(old_state.has_data); // spsc, nobody else is allowed to clear
            assert(self.consumer_slot != old_state.spare_slot);
            self.consumer_slot = old_state.spare_slot;
            self.consumer_done = false;
            return &self.bufs[self.consumer_slot.toIndex()];
        }

        pub fn consumerCommit(self: *@This()) void {
            self.consumer_done = true;
        }

        /// consumer only, non-blocking lock-free
        pub fn get(self: *@This()) ?T {
            const ptr = self.consumerPtr() orelse return null;
            const value = ptr.*;
            self.consumerCommit();
            return value;
        }
    };
}

test SPSCTripleBuffer {
    var tb: SPSCTripleBuffer(u8) = .{};
    try std.testing.expectEqual(null, tb.get());
    try std.testing.expectEqual(null, tb.get());
    try std.testing.expectEqual(null, tb.get());

    tb.put(3);
    try std.testing.expectEqual(3, tb.get());
    try std.testing.expectEqual(null, tb.get());
    try std.testing.expectEqual(null, tb.get());

    tb.put(4);
    tb.put(5);
    try std.testing.expectEqual(5, tb.get());
    try std.testing.expectEqual(null, tb.get());

    tb.put(6);
    try std.testing.expectEqual(6, tb.get());
    tb.put(7);
    try std.testing.expectEqual(7, tb.get());
    tb.put(8);
    try std.testing.expectEqual(8, tb.get());
    try std.testing.expectEqual(null, tb.get());

    tb.put(8);
    try std.testing.expectEqual(8, tb.get());
    tb.put(8);
    try std.testing.expectEqual(8, tb.get());
    tb.put(8);
    tb.put(8);
    tb.put(8);
    try std.testing.expectEqual(8, tb.get());
    try std.testing.expectEqual(null, tb.get());
}

test "SPSCTripleBuffer.init" {
    const tb: SPSCTripleBuffer(u8) = .{};
    try std.testing.expect(tb.consumer_slot != tb.producer_slot);
    try std.testing.expect(tb.consumer_slot != tb.shared_state.load(.seq_cst).spare_slot);
    try std.testing.expect(tb.producer_slot != tb.shared_state.load(.seq_cst).spare_slot);
}

test "SPSCTripleBuffer.producerPtr" {
    var tb: SPSCTripleBuffer(u8) = .{};
    try std.testing.expectEqual(null, tb.get());

    try std.testing.expectEqual(tb.producerPtr(), tb.producerPtr());

    tb.producerPtr().* = 3;
    try std.testing.expectEqual(null, tb.get());
    try std.testing.expectEqual(null, tb.get());
    tb.producerCommit();
    try std.testing.expectEqual(3, tb.get());
    try std.testing.expectEqual(null, tb.get());
}

test "SPSCTripleBuffer.consumerPtr" {
    var tb: SPSCTripleBuffer(u8) = .{};
    try std.testing.expectEqual(null, tb.get());

    try std.testing.expectEqual(null, tb.consumerPtr());

    tb.put(3);
    try std.testing.expectEqual(tb.consumerPtr(), tb.consumerPtr());
    try std.testing.expectEqual(tb.consumerPtr().?.*, tb.consumerPtr().?.*);
    tb.consumerCommit();
    try std.testing.expectEqual(null, tb.consumerPtr());
    try std.testing.expectEqual(null, tb.consumerPtr());
}

2 Likes

Awesome. Doing it lock-free is hard.

I wouldn’t mind adapting this to some apps I have that use mailboxes (currently wrapped in mutexes) .. but I absolutely need the reader functions to apply timeouts, so the idioms for the read loops look like

While forever, switch get next message or timeout
MsgA => processA(),
MsgB => processB(),
Timeout => handleIdleTimeout()

Any ideas on that ?

you can pair it with a std.Io.Event where you call set() from the producer after put and waitTimeout and reset from the consumer before get().