Batch bench

bench_test.go

package tidb

import (
	"math/rand"
	"testing"
)

var (
	totalCount = int64(1000000)
	batchSize  = int64(2000)
)

type Executor interface {
	Next() (int64, bool)
	NextBatch() ([]int64, bool)
}

type dataSource struct {
	currentCnt int64
	r          *rand.Rand
}

func (ds *dataSource) Next() (int64, bool) {
	if ds.currentCnt >= totalCount {
		return 0, false
	}
	ds.currentCnt += 1
	return ds.r.Int63(), true
}

func (ds *dataSource) NextBatch() ([]int64, bool) {
	if ds.currentCnt >= totalCount {
		return nil, false
	}
	var i int64
	res := make([]int64, 0, batchSize)
	for ; i < batchSize; i++ {
		res = append(res, ds.r.Int63())
	}
	ds.currentCnt += i
	return res, true
}

type dummyExecutor struct {
	src Executor
	r   *rand.Rand
}

func (de *dummyExecutor) Next() (int64, bool) {
	de.r.Int63()
	return de.src.Next()
}

func (de *dummyExecutor) NextBatch() ([]int64, bool) {
	rs, hasMore := de.src.NextBatch()
	if !hasMore {
		return nil, false
	}
	for _ = range rs {
		de.r.Int63()
	}
	return rs, hasMore
}

func composeExecutors() Executor {
	ds := &dataSource{r: rand.New(rand.NewSource(99))}
	de1 := &dummyExecutor{src: ds, r: rand.New(rand.NewSource(100))}
	de2 := &dummyExecutor{src: de1, r: rand.New(rand.NewSource(100))}
	de3 := &dummyExecutor{src: de2, r: rand.New(rand.NewSource(100))}
	de4 := &dummyExecutor{src: de3, r: rand.New(rand.NewSource(100))}
	de5 := &dummyExecutor{src: de4, r: rand.New(rand.NewSource(100))}
	return de5
}

func BenchmarkSingle(b *testing.B) {
	for i := 0; i < b.N; i++ {
		exec := composeExecutors()
		for {
			_, hasMore := exec.Next()
			if !hasMore {
				break
			}
		}
	}
}

func BenchmarkBatch(b *testing.B) {
	for i := 0; i < b.N; i++ {
		exec := composeExecutors()
		for {
			_, hasMore := exec.NextBatch()
			if !hasMore {
				break
			}
		}
	}
}