FIFO lock-less SPSC queue (unity 2022+)

SPSCQueue.cs

using System;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;
using JetBrains.Annotations;
using Unity.Collections.LowLevel.Unsafe;

/// <summary>
/// Simple Single-Producer Single-Consumer (SPSC) FIFO queue with batching support.
///
/// Capacity must be a power of 2 greater than 1.
/// 
/// Treat by definition: 'Single-Producer-Single-Consumer'
/// </summary>
[MustDisposeResource]
public sealed unsafe class SPSCQueue<T> : IDisposable where T : unmanaged
{
    public sealed class FullException    : Exception { public FullException()    : base("Queue is full.") { } }
    public sealed class EmptyException   : Exception { public EmptyException()   : base("Queue is empty.") { } }
    public sealed class TooManyElements  : Exception { public TooManyElements()  : base("Batch exceeds available capacity.") { } }

    // 128 bytes covers both 64-byte (x64) and 128-byte (Apple Silicon) cache lines.
    [StructLayout(LayoutKind.Explicit, Size = 128)]
    private struct PaddedLong
    {
        [FieldOffset(0)]
        public long Value;
    }

    private PaddedLong _head;
    private PaddedLong _tail;

    private readonly T*  _buffer;
    private readonly int _capacity;
    private readonly int _mask;
    private bool         _disposed;

    /// <param name="capacity">Must be a power of 2 greater than 1.</param>
    public SPSCQueue(int capacity)
    {
        if (capacity <= 1 || (capacity & (capacity - 1)) != 0)
            throw new ArgumentException("Capacity must be a power of 2 greater than 1.", nameof(capacity));

        // 64 byte aligned, unmanaged heap.
        _buffer   = (T*)UnsafeUtility.Malloc(
                        size:      capacity * sizeof(T),
                        alignment: 64,
                        allocator: Unity.Collections.Allocator.Persistent);
        _capacity = capacity;
        _mask     = capacity - 1;

        Volatile.Write(ref _head.Value, 0L);
        Volatile.Write(ref _tail.Value, 0L);
    }

    public int Capacity => _capacity;

    // -------------------------------------------------------------------------
    // Single-element operations
    // -------------------------------------------------------------------------

    /// <summary>Producer only. Returns false if the queue is full instead of throwing.</summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool TryEnqueue(T value)
    {
        long tail = Volatile.Read(ref _tail.Value);
        long next = (tail + 1) & _mask;

        if (next == Volatile.Read(ref _head.Value))
            return false;

        _buffer[tail] = value;
        Volatile.Write(ref _tail.Value, next);
        return true;
    }
    
    /// <summary>Consumer only. Returns the head element without removing it.
    /// Returns false the queue is empty.</summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool TryPeek(out T result)
    {
        long head = Volatile.Read(ref _head.Value);
        long tail = Volatile.Read(ref _tail.Value);   // ACQUIRE

        if (head == tail)
        {
            result = default;
            return false;
        }
        
        result = _buffer[head];
        return true;
    }

    /// <summary>Consumer only. Returns false if the queue is empty instead of throwing.</summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool TryDequeue(out T value)
    {
        long head = Volatile.Read(ref _head.Value);

        if (head == Volatile.Read(ref _tail.Value))
        {
            value = default;
            return false;
        }

        value     = _buffer[head];
        long next = (head + 1) & _mask;

        Volatile.Write(ref _head.Value, next);
        return true;
    }

    // -------------------------------------------------------------------------
    // Batched operations
    // -------------------------------------------------------------------------

    /// <summary>
    /// Producer only. Enqueues all items atomically (single RELEASE store after all writes).
    /// Returns false if the batch won't fit.
    /// </summary>
    public bool EnqueueBatch(ReadOnlySpan<T> items)
    {
        int  len       = _capacity;
        long mask      = _mask;
        long tail      = Volatile.Read(ref _tail.Value);
        long head      = Volatile.Read(ref _head.Value);   // ACQUIRE
        long used      = (tail - head + len) & mask;
        long available = len - 1 - used;                   // -1 for sentinel slot

        if (available < items.Length)
            return false;

        int count   = items.Length;
        int tailIdx = (int)tail;
        int tillEnd = len - tailIdx;

        if (count <= tillEnd)
        {
            items[..count].CopyTo(new Span<T>(_buffer + tailIdx, count));
        }
        else
        {
            items[..tillEnd].CopyTo(new Span<T>(_buffer + tailIdx, tillEnd));
            items[tillEnd..count].CopyTo(new Span<T>(_buffer, count - tillEnd));
        }

        Volatile.Write(ref _tail.Value, (tail + count) & mask);   // RELEASE
        return true;
    }

    /// <summary>
    /// Consumer only. Dequeues up to <c>out.Length</c> elements.
    /// Returns the number of elements actually dequeued.
    /// Returns 0 if the queue is empty.
    /// </summary>
    public int DequeueBatch(Span<T> @out)
    {
        int  len       = _capacity;
        long mask      = _mask;
        long head      = Volatile.Read(ref _head.Value);
        long tail      = Volatile.Read(ref _tail.Value);   // ACQUIRE

        if (head == tail) return 0;

        long available = (tail - head + len) & mask;
        int  count     = (int)Math.Min(available, @out.Length);
        int  headIdx   = (int)head;
        int  tillEnd   = len - headIdx;

        if (count <= tillEnd)
        {
            new ReadOnlySpan<T>(_buffer + headIdx, count).CopyTo(@out[..count]);
        }
        else
        {
            new ReadOnlySpan<T>(_buffer + headIdx, tillEnd).CopyTo(@out[..tillEnd]);
            new ReadOnlySpan<T>(_buffer, count - tillEnd).CopyTo(@out[tillEnd..count]);
        }

        Volatile.Write(ref _head.Value, (head + count) & mask);   // RELEASE
        return count;
    }

    // -------------------------------------------------------------------------
    // Lifetime
    // -------------------------------------------------------------------------

    public void Dispose()
    {
        if (_disposed) return;
        _disposed = true;
        UnsafeUtility.Free(_buffer, Unity.Collections.Allocator.Persistent);
    }
}