lock-free ringbuffer

ringbuffer.go

package main

import (
	"fmt"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
	"unsafe"
)

const N = 100_000
const K = 100_000_000

type pad64[T any] struct {
	_ [64]byte
	v T
}

type ringBuffer[T any] struct {
	buf       pad64[unsafe.Pointer]
	head      pad64[atomic.Uint64]
	headCache pad64[uint64]
	tail      pad64[atomic.Uint64]
	tailCache pad64[uint64]
}

//go:linkname mallocgc runtime.mallocgc
func mallocgc(size uintptr, typ unsafe.Pointer, needzero bool) unsafe.Pointer

func new[T any]() *ringBuffer[T] {
	var x T
	buf := mallocgc(uintptr(N)*unsafe.Sizeof(x), nil, false)
	return &ringBuffer[T]{
		buf: pad64[unsafe.Pointer]{v: buf},
	}
}

func (_ *ringBuffer[T]) next(curr uint64) uint64 {
	n := curr + 1
	if n == N {
		return 0
	}
	return n
}

func (rb *ringBuffer[T]) push(v T) bool {
	head := rb.head.v.Load()
	next := rb.next(head)

	if next == rb.tailCache.v {
		rb.tailCache.v = rb.tail.v.Load()
		if next == rb.tailCache.v {
			return false
		}
	}

	*(*T)(unsafe.Add(rb.buf.v, head*uint64(unsafe.Sizeof(v)))) = v
	rb.head.v.Store(next)

	return true
}

func (rb *ringBuffer[T]) pop(v *T) bool {
	tail := rb.tail.v.Load()

	if tail == rb.headCache.v {
		rb.headCache.v = rb.head.v.Load()
		if tail == rb.headCache.v {
			return false
		}
	}

	*v = *(*T)(unsafe.Add(rb.buf.v, tail*uint64(unsafe.Sizeof(*v))))
	rb.tail.v.Store(rb.next(tail))

	return true
}

func (rb *ringBuffer[T]) release() {
	// TOOD
}

func main() {
	runtime.LockOSThread()

	rb := new[int32]()

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		runtime.LockOSThread()
		var v int32
		for i := 0; i < K; i++ {
			for !rb.pop(&v) {
			}
			if v != int32(i) {
				panic("incorrect value")
			}
		}
		wg.Done()
	}()

	t := time.Now()
	for i := 0; i < K; i++ {
		for !rb.push(int32(i)) {
		}
	}
	wg.Wait()
	elapsed := time.Since(t)

	fmt.Printf("elapsed %d ms, %d ns/op, %d op/sec\n", elapsed.Milliseconds(), uint64(elapsed.Nanoseconds()/K), uint64(float64(K)/elapsed.Seconds()))
}