lotsaa.go

package lotsa

import (
	"fmt"
	"io"
	"math/rand"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

// Output is used to print elapsed time and ops/sec
var Output io.Writer

// MemUsage is used to output the memory usage
var MemUsage bool

// Ops executed a number of operations over a multiple goroutines.
// count is the number of operations.
// threads is the number goroutines.
// op is the operation function
func Ops(count, threads int, op func(i, thread int)) {
	var start time.Time
	var wg sync.WaitGroup
	wg.Add(threads)
	var ms1 runtime.MemStats
	output := Output
	if output != nil {
		if MemUsage {
			runtime.GC()
			runtime.ReadMemStats(&ms1)
		}
		start = time.Now()
	}
	for i := 0; i < threads; i++ {
		s, e := count/threads*i, count/threads*(i+1)
		if i == threads-1 {
			e = count
		}
		go func(i, s, e int) {
			defer wg.Done()
			for j := s; j < e; j++ {
				op(j, i)
			}
		}(i, s, e)
	}
	wg.Wait()

	if output != nil {
		dur := time.Since(start)
		var alloc uint64
		if MemUsage {
			runtime.GC()
			var ms2 runtime.MemStats
			runtime.ReadMemStats(&ms2)
			if ms1.HeapAlloc > ms2.HeapAlloc {
				alloc = 0
			} else {
				alloc = ms2.HeapAlloc - ms1.HeapAlloc
			}
		}
		WriteOutput(output, int64(count), threads, dur, alloc)
	}
}

// Time executed operations over multiple goroutines for a fixed duration.
// duration is the duration for continuously running the Operation.
// threads is the number goroutines.
// op is the operation function
func Time(duration time.Duration, threads int, op func(threadRand *rand.Rand, thread int)) {

	var start time.Time
	var wg sync.WaitGroup
	wg.Add(threads)
	var ms1 runtime.MemStats
	output := Output
	if output != nil {
		if MemUsage {
			runtime.GC()
			runtime.ReadMemStats(&ms1)
		}
		start = time.Now()
	}

	var count int64
	for i := 0; i < threads; i++ {

		go func(i int) {
			threadRand := rand.New(rand.NewSource(time.Now().UnixNano()))
			done := time.After(duration)
			var localCount int64 = 0
			for {
				select {
				case <-done:
					wg.Done()
					atomic.AddInt64(&count, localCount)
					return
				default:
					{
						op(threadRand, i)
						localCount++
					}
				}
			}
		}(i)
	}
	wg.Wait()

	if output != nil {
		dur := time.Since(start)
		var alloc uint64
		if MemUsage {
			runtime.GC()
			var ms2 runtime.MemStats
			runtime.ReadMemStats(&ms2)
			if ms1.HeapAlloc > ms2.HeapAlloc {
				alloc = 0
			} else {
				alloc = ms2.HeapAlloc - ms1.HeapAlloc
			}
		}
		WriteOutput(output, count, threads, dur, alloc)
	}
}

func commaize(n int64) string {
	s1, s2 := fmt.Sprintf("%d", n), ""
	for i, j := len(s1)-1, 0; i >= 0; i, j = i-1, j+1 {
		if j%3 == 0 && j != 0 {
			s2 = "," + s2
		}
		s2 = string(s1[i]) + s2
	}
	return s2
}

func memstr(alloc uint64) string {
	switch {
	case alloc <= 1024:
		return fmt.Sprintf("%d bytes", alloc)
	case alloc <= 1024*1024:
		return fmt.Sprintf("%.1f KB", float64(alloc)/1024)
	case alloc <= 1024*1024*1024:
		return fmt.Sprintf("%.1f MB", float64(alloc)/1024/1024)
	default:
		return fmt.Sprintf("%.1f GB", float64(alloc)/1024/1024/1024)
	}
}

// WriteOutput writes an output line to the specified writer
func WriteOutput(w io.Writer, count int64, threads int, elapsed time.Duration, alloc uint64) {
	var ss string
	if threads != 1 {
		ss = fmt.Sprintf("over %d threads ", threads)
	}
	var nsop int
	if count > 0 {
		nsop = int(elapsed / time.Duration(count))
	}
	var allocstr string
	if alloc > 0 {
		var bops float64
		if count > 0 {
			bops = float64(alloc) / float64(count)
		}
		allocstr = fmt.Sprintf(", %s, %.1f bytes/op", memstr(alloc), bops)
	}
	fmt.Fprintf(w, "%s ops %sin %.0fms, %s/sec, %d ns/op%s\n",
		commaize(count), ss, elapsed.Seconds()*1000,
		commaize(int64(float64(count)/elapsed.Seconds())),
		nsop, allocstr,
	)
}