internal/ringbuffer/ringbuffer.go

// Package ringbuffer provides lock-free ring buffer implementations for high-performance streaming.
// These are used for stream chunk queuing with minimal synchronization overhead.
package ringbuffer

import (
	"runtime"
	"sync/atomic"
)

const (
	// CacheLineSize is the typical CPU cache line size (64 bytes on most architectures)
	CacheLineSize = 64
)

// RingBuffer is a single-producer single-consumer (SPSC) lock-free ring buffer.
// It uses atomic operations for synchronization, making it suitable for high-throughput
// scenarios where minimal latency is critical.
type RingBuffer[T any] struct {
	// Padding to prevent false sharing between head and tail
	_ [CacheLineSize]byte

	// head is the read position (consumer)
	head atomic.Uint64

	_ [CacheLineSize - 8]byte // Padding to separate head and tail to different cache lines

	// tail is the write position (producer)
	tail atomic.Uint64

	_ [CacheLineSize - 8]byte // Padding

	buffer []T
	mask   uint64
}

// New creates a new RingBuffer with capacity rounded up to power of 2.
func New[T any](capacity int) *RingBuffer[T] {
	if capacity <= 0 {
		capacity = 1024
	}

	// Round up to power of 2
	capacity = roundUpPowerOf2(capacity)

	return &RingBuffer[T]{
		buffer: make([]T, capacity),
		mask:   uint64(capacity - 1),
	}
}

// Push adds an item to the ring buffer.
// Returns false if the buffer is full.
// This method is safe for single-producer usage.
func (r *RingBuffer[T]) Push(item T) bool {
	tail := r.tail.Load()
	head := r.head.Load()

	// Check if full
	if tail-head >= uint64(len(r.buffer)) {
		return false
	}

	// Write item
	r.buffer[tail&r.mask] = item

	// Update tail with memory barrier
	r.tail.Store(tail + 1)

	return true
}

// Pop removes and returns an item from the ring buffer.
// Returns (zero value, false) if the buffer is empty.
// This method is safe for single-consumer usage.
func (r *RingBuffer[T]) Pop() (T, bool) {
	var zero T

	head := r.head.Load()
	tail := r.tail.Load()

	// Check if empty
	if head == tail {
		return zero, false
	}

	// Read item
	item := r.buffer[head&r.mask]

	// Update head with memory barrier
	r.head.Store(head + 1)

	return item, true
}

// TryPush attempts to add an item with a spin-wait (for batch operations).
func (r *RingBuffer[T]) TryPush(item T, spins int) bool {
	for i := 0; i < spins; i++ {
		if r.Push(item) {
			return true
		}
		runtime.Gosched()
	}
	return false
}

// TryPop attempts to remove an item with a spin-wait.
func (r *RingBuffer[T]) TryPop(spins int) (T, bool) {
	var zero T
	for i := 0; i < spins; i++ {
		if item, ok := r.Pop(); ok {
			return item, true
		}
		runtime.Gosched()
	}
	return zero, false
}

// Len returns the current number of items in the buffer.
func (r *RingBuffer[T]) Len() int {
	return int(r.tail.Load() - r.head.Load())
}

// Cap returns the capacity of the buffer.
func (r *RingBuffer[T]) Cap() int {
	return len(r.buffer)
}

// IsEmpty returns true if the buffer is empty.
func (r *RingBuffer[T]) IsEmpty() bool {
	return r.head.Load() == r.tail.Load()
}

// IsFull returns true if the buffer is full.
func (r *RingBuffer[T]) IsFull() bool {
	return r.tail.Load()-r.head.Load() >= uint64(len(r.buffer))
}

// Reset clears the buffer (not thread-safe).
func (r *RingBuffer[T]) Reset() {
	r.head.Store(0)
	r.tail.Store(0)
	for i := range r.buffer {
		var zero T
		r.buffer[i] = zero
	}
}

// MPMCQueue is a multi-producer multi-consumer lock-free queue.
// It uses a different algorithm that supports concurrent access from multiple goroutines.
type MPMCQueue[T any] struct {
	_     [CacheLineSize]byte
	slots []slot[T]
	mask  uint64
	sendx atomic.Uint64
	recvx atomic.Uint64
	_     [CacheLineSize]byte
}

type slot[T any] struct {
	turn atomic.Uint64
	item T
	_    [CacheLineSize - 16]byte
}

// NewMPMC creates a new multi-producer multi-consumer queue.
func NewMPMC[T any](capacity int) *MPMCQueue[T] {
	if capacity <= 0 {
		capacity = 1024
	}
	capacity = roundUpPowerOf2(capacity)

	q := &MPMCQueue[T]{
		slots: make([]slot[T], capacity),
		mask:  uint64(capacity - 1),
	}

	// Initialize turn counters
	for i := range q.slots {
		q.slots[i].turn.Store(uint64(i))
	}

	return q
}

// Enqueue adds an item to the queue.
func (q *MPMCQueue[T]) Enqueue(item T) bool {
	for {
		sendx := q.sendx.Load()
		slot := &q.slots[sendx&q.mask]
		turn := slot.turn.Load()

		if turn == sendx {
			if q.sendx.CompareAndSwap(sendx, sendx+1) {
				slot.item = item
				slot.turn.Store(sendx + 1)
				return true
			}
		} else if turn < sendx {
			// Queue is full
			return false
		}
		// Slot not ready, retry
		runtime.Gosched()
	}
}

// Dequeue removes and returns an item from the queue.
func (q *MPMCQueue[T]) Dequeue() (T, bool) {
	var zero T

	for {
		recvx := q.recvx.Load()
		slot := &q.slots[recvx&q.mask]
		turn := slot.turn.Load()

		if turn == recvx+1 {
			if q.recvx.CompareAndSwap(recvx, recvx+1) {
				item := slot.item
				var empty T
				slot.item = empty
				slot.turn.Store(recvx + q.mask + 1)
				return item, true
			}
		} else if turn < recvx+1 {
			// Queue is empty
			return zero, false
		}
		// Slot not ready, retry
		runtime.Gosched()
	}
}

// Len returns approximate queue length.
func (q *MPMCQueue[T]) Len() int {
	return int(q.sendx.Load() - q.recvx.Load())
}

func roundUpPowerOf2(n int) int {
	if n <= 0 {
		return 1
	}
	if n&(n-1) == 0 {
		return n
	}
	p := 1
	for p < n {
		p <<= 1
	}
	return p
}