trbbench.go

package main

import (
	"flag"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"runtime"
	"time"

	"github.com/cespare/wait"
	"github.com/liftoffio/foundationdb/bindings/go/src/fdb"
	"golang.org/x/sys/unix"
)

func randByteSlice(r *rand.Rand, n int) []byte {
	b := make([]byte, n)
	for i := 0; i < n; i++ {
		b[i] = byte(r.Intn(255))
	}
	return b
}

type kk struct {
	k []byte
}

func (kk kk) FDBKey() fdb.Key {
	return kk.k
}

func main() {
	var (
		numWorkers  = flag.Int("workers", 1000, "Number of workers")
		clusterFile = flag.String("clusterfile", "fdb.cluster", "FoundationDB cluster file")
	)
	flag.Parse()
	if err := fdb.APIVersion(600); err != nil {
		log.Fatalf("error setting API version: %s", err)
	}
	db, err := fdb.Open(*clusterFile, []byte("DB"))
	if err != nil {
		log.Fatal(err)
	}

	quit := make(chan struct{})
	nowUnix := time.Now().Unix()

	var wg wait.Group
	for i := 0; i < *numWorkers; i++ {
		r := rand.New(rand.NewSource(nowUnix - int64(i+1)))
		wg.Go(func(<-chan struct{}) error {
			for {
				select {
				case <-quit:
					return nil
				default:
				}
				kk := kk{randByteSlice(r, 50)}
				v := randByteSlice(r, 7000)
				// put
				_, err := db.Transact(func(tr fdb.Transaction) (i interface{}, e error) {
					tr.Set(kk, v)
					return nil, nil
				})
				if err != nil {
					log.Println(err)
					return err
				}
			}
			return nil
		})
	}

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, unix.SIGQUIT)
	go func() {
		for range ch {
			os.Stderr.Write(GoroutineStacks())
		}
	}()

	go func() {
		defer close(quit)
		signals := make(chan os.Signal, 1)
		// SIGTERM and SIGINT are used by runit and humans, respectively, to ask
		// the process to stop.
		signal.Notify(signals, unix.SIGINT, unix.SIGTERM)
		<-signals
	}()

	if err := wg.Wait(); err != nil {
		log.Fatal(err)
	}
}

// GoroutineStacks returns goroutine stack traces for all running goroutines.
// (It is like runtime/debug.Stack, but for all goroutines rather than just the
// calling goroutine.)
func GoroutineStacks() []byte {
	// We cannot know the required buffer size, so start with 1MB and double
	// up to a maximum of 64MB before giving up and using a truncated trace.
	var buf []byte
	for size := int(1e6); size <= 64e6; size *= 2 {
		buf = make([]byte, size)
		n := runtime.Stack(buf, true)
		if n < len(buf) {
			buf = buf[:n]
			break
		}
	}
	return buf
}