Merge pull request #6 from asymmetric-research/marctrem/make-concurrent-writing-to-ringbuf-racefree

make concurrent writing to ringbuf racefree

Author: nlgripto

collections/ringbuffer/ringbuffer.go

@@ -2,16 +2,17 @@ package ringbuffer
 
 import (
 	"fmt"
+	"sync/atomic"
 
 	armath "github.com/asymmetric-research/go-commons/math"
 )
 
 type T[C any] struct {
-	buflen uint
+	buflen uint64
 	buf    []C
 
 	// head points to the next free slot
-	head uint
+	head atomic.Uint64
 }
 
 func New[C any](size int) (*T[C], error) {
@@ -26,29 +27,29 @@ func NewInto[C any](dst *T[C], buf []C) error {
 		return fmt.Errorf("backing buffer must have a greater than zero")
 	}
 	*dst = T[C]{
-		buflen: uint(len(buf)),
+		buflen: uint64(len(buf)),
 		buf:    buf,
-		head:   0,
 	}
+	dst.head.Store(0)
 	return nil
 }
 
 func (r *T[C]) Push(item C) {
-	r.buf[r.head%r.buflen] = item
-	r.head += 1
+	nextSlot := r.head.Add(1)
+	r.buf[(nextSlot-1)%r.buflen] = item
 }
 
 func (r *T[C]) Last(dst []C) int {
 	// how many entries can we write?
-	maxWritable := armath.Min(r.head, r.buflen)
+	maxWritable := armath.Min(r.head.Load(), r.buflen)
 
 	// if the dst is larger than the amount of entries we can write, let's clamp it.
 	if len(dst) > int(maxWritable) {
 		// only consider the first available slots of dst
 		dst = dst[:maxWritable]
 	}
 
-	headmod := int(r.head % r.buflen)
+	headmod := int(r.head.Load() % r.buflen)
 
 	// we must do at most 2 copies
 	n := 0
@@ -80,8 +81,8 @@ func (r *T[C]) Last(dst []C) int {
 	return n
 }
 
-func (r *T[C]) Len() uint {
-	used := armath.Min(r.buflen, r.head)
+func (r *T[C]) Len() uint64 {
+	used := armath.Min(r.buflen, r.head.Load())
 	return used
 }
 
@@ -92,17 +93,19 @@ const (
 	SEQ_MODE_FILO
 )
 
-func (r *T[C]) Seq(seqMode SeqMode) func(yield func(uint, C) bool) {
-	return func(yield func(uint, C) bool) {
+func (r *T[C]) Seq(seqMode SeqMode) func(yield func(uint64, C) bool) {
+	return func(yield func(uint64, C) bool) {
 		if r.buflen == 0 {
 			return
 		}
 
+		head := r.head.Load()
+
 		// how many entries can we write?
-		maxWritable := armath.Min(r.head, r.buflen)
+		maxWritable := armath.Min(head, r.buflen)
 
 		if seqMode == SEQ_MODE_FIFO {
-			start := (((r.head - 1) % r.buflen) - maxWritable) % r.buflen
+			start := (((head - 1) % r.buflen) - maxWritable) % r.buflen
 
 			for i := range maxWritable {
 				idx := (start + i) % r.buflen
@@ -113,7 +116,7 @@ func (r *T[C]) Seq(seqMode SeqMode) func(yield func(uint, C) bool) {
 			return
 		}
 		if seqMode == SEQ_MODE_FILO {
-			start := r.head - 1
+			start := head - 1
 			for i := range maxWritable {
 				idx := (start - i) % r.buflen
 				if !yield(i, r.buf[idx]) {