benchmark_busy.go

package busy

import (
	"runtime"
	"sync/atomic"
	"testing"
)

type Queue interface {
	EnqueueRequest()
	NextRequestToSend() bool
}

const size = 64
const mask = size - 1

type ring struct {
	_     [8]uint64
	write uint64
	_     [7]uint64
	read1 uint64
	_     [7]uint64
	slots [size]slot
}

type slot struct {
	mark uint32
}

func (r *ring) EnqueueRequest() {
	s := &r.slots[atomic.AddUint64(&r.write, 1)&mask]
	for !atomic.CompareAndSwapUint32(&s.mark, 0, 1) {
		runtime.Gosched()
	}
}

func (r *ring) NextRequestToSend() (ret bool) {
	r.read1++
	s := &r.slots[r.read1&mask]
	if ret = atomic.CompareAndSwapUint32(&s.mark, 1, 0); !ret {
		r.read1--
	}
	return ret
}

type channel struct {
	ch chan struct{}
}

func (c *channel) EnqueueRequest() {
	c.ch <- struct{}{}
}

func (c *channel) NextRequestToSend() bool {
	select {
	case <-c.ch:
		return true
	default:
		return false
	}
}

func BenchmarkBusy(b *testing.B) {
	bench := func(name string, parallelism int, mk func() Queue) {
		b.Run(name, func(b *testing.B) {
			q := mk()
			s := uint32(0)
			done := make(chan struct{})
			go func() {
				for atomic.LoadUint32(&s) == 0 {
					if !q.NextRequestToSend() {
						runtime.Gosched()
					}
				}
				close(done)
			}()
			b.ResetTimer()
			b.SetParallelism(parallelism)
			b.RunParallel(func(pb *testing.PB) {
				for pb.Next() {
					q.EnqueueRequest()
				}
			})
			b.StopTimer()
			if c, ok := q.(*channel); ok {
				close(c.ch)
			}
			atomic.StoreUint32(&s, 1)
			<-done
		})
	}
	bench("Chan-size-64-parallelism-01", 1, func() Queue { return &channel{ch: make(chan struct{}, size)} })
	bench("Ring-size-64-parallelism-01", 1, func() Queue { return &ring{} })

	bench("Chan-size-64-parallelism-08", 8, func() Queue { return &channel{ch: make(chan struct{}, size)} })
	bench("Ring-size-64-parallelism-08", 8, func() Queue { return &ring{} })

	bench("Chan-size-64-parallelism-64", 64, func() Queue { return &channel{ch: make(chan struct{}, size)} })
	bench("Ring-size-64-parallelism-64", 64, func() Queue { return &ring{} })
}