382 lines
9.1 KiB
Go
382 lines
9.1 KiB
Go
package gen
|
|
|
|
import (
|
|
"encoding"
|
|
"encoding/binary"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"iter"
|
|
"math/rand/v2"
|
|
"os"
|
|
|
|
"github.com/cespare/xxhash/v2"
|
|
|
|
module "cosmossdk.io/api/cosmos/benchmark/module/v1"
|
|
"cosmossdk.io/tools/benchmark"
|
|
)
|
|
|
|
// Options is the configuration for the generator.
|
|
type Options struct {
|
|
*module.GeneratorParams
|
|
// HomeDir is for reading/writing state
|
|
HomeDir string
|
|
|
|
InsertWeight float64
|
|
UpdateWeight float64
|
|
GetWeight float64
|
|
DeleteWeight float64
|
|
}
|
|
|
|
// State is the state of the generator.
|
|
// It can be marshaled and unmarshaled to/from a binary format.
|
|
type State struct {
|
|
Src interface {
|
|
rand.Source
|
|
encoding.BinaryMarshaler
|
|
encoding.BinaryUnmarshaler
|
|
}
|
|
Keys [][]Payload
|
|
}
|
|
|
|
// Marshal writes the state to w.
|
|
func (s *State) Marshal(w io.Writer) error {
|
|
srcBz, err := s.Src.MarshalBinary()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var n int
|
|
n, err = w.Write(srcBz)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if n != 20 {
|
|
return fmt.Errorf("expected 20 bytes, got %d", n)
|
|
}
|
|
if err = binary.Write(w, binary.LittleEndian, uint64(len(s.Keys))); err != nil {
|
|
return err
|
|
}
|
|
for _, bucket := range s.Keys {
|
|
if err = binary.Write(w, binary.LittleEndian, uint64(len(bucket))); err != nil {
|
|
return err
|
|
}
|
|
for _, key := range bucket {
|
|
if err = binary.Write(w, binary.LittleEndian, key); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Unmarshal reads the state from r.
|
|
func (s *State) Unmarshal(r io.Reader) error {
|
|
srcBz := make([]byte, 20)
|
|
if _, err := r.Read(srcBz); err != nil {
|
|
return err
|
|
}
|
|
s.Src = rand.NewPCG(0, 0)
|
|
if err := s.Src.UnmarshalBinary(srcBz); err != nil {
|
|
return err
|
|
}
|
|
|
|
var n uint64
|
|
if err := binary.Read(r, binary.LittleEndian, &n); err != nil {
|
|
return err
|
|
}
|
|
s.Keys = make([][]Payload, n)
|
|
for i := uint64(0); i < n; i++ {
|
|
var m uint64
|
|
if err := binary.Read(r, binary.LittleEndian, &m); err != nil {
|
|
return err
|
|
}
|
|
s.Keys[i] = make([]Payload, m)
|
|
for j := uint64(0); j < m; j++ {
|
|
if err := binary.Read(r, binary.LittleEndian, &s.Keys[i][j]); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Generator generates operations for a benchmark transaction.
|
|
// The generator is stateful, keeping track of which keys have been inserted
|
|
// so that meaningful gets and deletes can be generated.
|
|
type Generator struct {
|
|
Options
|
|
|
|
rand *rand.Rand
|
|
state *State
|
|
}
|
|
|
|
type opt func(*Generator)
|
|
|
|
// NewGenerator creates a new generator with the given options.
|
|
func NewGenerator(opts Options, f ...opt) *Generator {
|
|
g := &Generator{
|
|
Options: opts,
|
|
state: &State{
|
|
Src: rand.NewPCG(opts.Seed, opts.Seed>>32),
|
|
},
|
|
}
|
|
g.rand = rand.New(g.state.Src)
|
|
for _, fn := range f {
|
|
fn(g)
|
|
}
|
|
return g
|
|
}
|
|
|
|
// WithGenesis sets the generator state to the genesis seed.
|
|
// When the generator is created, it will sync to genesis state.
|
|
// The benchmark client needs to do this so that it can generate meaningful tx operations.
|
|
func WithGenesis() func(*Generator) {
|
|
return func(g *Generator) {
|
|
// sync state to genesis seed
|
|
g.state.Keys = make([][]Payload, g.BucketCount)
|
|
if g.GeneratorParams != nil {
|
|
for kv := range g.GenesisSet() {
|
|
g.state.Keys[kv.StoreKey] = append(g.state.Keys[kv.StoreKey], kv.Key)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// WithSeed sets the seed for the generator.
|
|
func WithSeed(seed uint64) func(*Generator) {
|
|
return func(g *Generator) {
|
|
g.state.Src = rand.NewPCG(seed, seed>>32)
|
|
g.rand = rand.New(g.state.Src)
|
|
}
|
|
}
|
|
|
|
// Load loads the generator state from disk.
|
|
func (g *Generator) Load() error {
|
|
f := fmt.Sprintf("%s/data/generator_state.bin", g.HomeDir)
|
|
r, err := os.Open(f)
|
|
if err != nil {
|
|
if errors.Is(err, os.ErrNotExist) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
return g.state.Unmarshal(r)
|
|
}
|
|
|
|
// Payload is a 2-tuple of seed and length.
|
|
// A seed is uint64 which is used to generate a byte slice of size length.
|
|
type Payload [2]uint64
|
|
|
|
// Seed returns the seed in the payload.
|
|
func (p Payload) Seed() uint64 {
|
|
return p[0]
|
|
}
|
|
|
|
// Length returns the length in the payload.
|
|
func (p Payload) Length() uint64 {
|
|
return p[1]
|
|
}
|
|
|
|
// Bytes returns the byte slice generated from the seed and length.
|
|
// The underlying byte slice is deterministically generated using the (very fast) xxhash algorithm.
|
|
func (p Payload) Bytes() []byte {
|
|
return Bytes(p.Seed(), p.Length())
|
|
}
|
|
|
|
func (p Payload) String() string {
|
|
return fmt.Sprintf("(%d, %d)", p.Seed(), p.Length())
|
|
}
|
|
|
|
func NewPayload(seed, length uint64) Payload {
|
|
return Payload{seed, length}
|
|
}
|
|
|
|
// KV is a key-value pair with a store key.
|
|
type KV struct {
|
|
StoreKey uint64
|
|
Key Payload
|
|
Value Payload
|
|
}
|
|
|
|
func (g *Generator) fetchKey(bucket uint64) (idx uint64, key Payload, err error) {
|
|
bucketLen := uint64(len(g.state.Keys[bucket]))
|
|
if bucketLen == 0 {
|
|
return 0, Payload{}, fmt.Errorf("no keys in bucket %d", bucket)
|
|
}
|
|
idx = g.rand.Uint64N(bucketLen)
|
|
return idx, g.state.Keys[bucket][idx], nil
|
|
}
|
|
|
|
func (g *Generator) deleteKey(bucket, idx uint64) {
|
|
g.state.Keys[bucket] = append(g.state.Keys[bucket][:idx], g.state.Keys[bucket][idx+1:]...)
|
|
}
|
|
|
|
func (g *Generator) setKey(bucket uint64, payload Payload) {
|
|
g.state.Keys[bucket] = append(g.state.Keys[bucket], payload)
|
|
}
|
|
|
|
// GenesisSet returns a sequence of key-value pairs for the genesis state.
|
|
// It is called by the server during InitGenesis to generate and set the initial state.
|
|
// The client uses WithGenesis to sync to the genesis state.
|
|
func (g *Generator) GenesisSet() iter.Seq[*KV] {
|
|
return func(yield func(*KV) bool) {
|
|
for range g.GenesisCount {
|
|
seed := g.rand.Uint64()
|
|
if !yield(&KV{
|
|
StoreKey: g.UintN(g.BucketCount),
|
|
Key: NewPayload(seed, g.getLength(g.KeyMean, g.KeyStdDev)),
|
|
Value: NewPayload(seed, g.getLength(g.ValueMean, g.ValueStdDev)),
|
|
}) {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Next generates the next benchmark operation.
|
|
// The operation is one of insert, update, get, or delete.
|
|
// The tx client calls this function to deterministically generate the next operation.
|
|
func (g *Generator) Next() (uint64, *benchmark.Op, error) {
|
|
if g.InsertWeight+g.UpdateWeight+g.GetWeight+g.DeleteWeight != 1 {
|
|
return 0, nil, fmt.Errorf("weights must sum to 1")
|
|
}
|
|
|
|
var (
|
|
err error
|
|
key Payload
|
|
)
|
|
x := g.rand.Float64()
|
|
bucket := g.UintN(g.BucketCount)
|
|
op := &benchmark.Op{
|
|
Exists: true,
|
|
}
|
|
|
|
switch {
|
|
case x < g.InsertWeight:
|
|
// insert
|
|
op.Seed = g.rand.Uint64()
|
|
op.KeyLength = g.getLength(g.KeyMean, g.KeyStdDev)
|
|
op.ValueLength = g.getLength(g.ValueMean, g.ValueStdDev)
|
|
op.Exists = false
|
|
g.setKey(bucket, NewPayload(op.Seed, op.KeyLength))
|
|
case x < g.InsertWeight+g.UpdateWeight:
|
|
// update
|
|
_, key, err = g.fetchKey(bucket)
|
|
if err != nil {
|
|
return 0, nil, err
|
|
}
|
|
op.Seed = key.Seed()
|
|
op.KeyLength = key.Length()
|
|
op.ValueLength = g.getLength(g.ValueMean, g.ValueStdDev)
|
|
case x < g.InsertWeight+g.UpdateWeight+g.GetWeight:
|
|
// get
|
|
_, key, err = g.fetchKey(bucket)
|
|
if err != nil {
|
|
return 0, nil, err
|
|
}
|
|
op.Seed = key.Seed()
|
|
op.KeyLength = key.Length()
|
|
default:
|
|
// delete
|
|
var idx uint64
|
|
idx, key, err = g.fetchKey(bucket)
|
|
if err != nil {
|
|
return 0, nil, err
|
|
}
|
|
op.Delete = true
|
|
op.Seed = key.Seed()
|
|
op.KeyLength = key.Length()
|
|
g.deleteKey(bucket, idx)
|
|
}
|
|
|
|
return bucket, op, nil
|
|
}
|
|
|
|
// NormUint64 returns a random uint64 with a normal distribution.
|
|
func (g *Generator) NormUint64(mean, stdDev uint64) uint64 {
|
|
return uint64(g.rand.NormFloat64()*float64(stdDev) + float64(mean))
|
|
}
|
|
|
|
func (g *Generator) getLength(mean, stdDev uint64) uint64 {
|
|
length := g.NormUint64(mean, stdDev)
|
|
if length == 0 {
|
|
length = 1
|
|
}
|
|
return length
|
|
}
|
|
|
|
// UintN returns a random uint64 in the range [0, n).
|
|
func (g *Generator) UintN(n uint64) uint64 {
|
|
return g.rand.Uint64N(n)
|
|
}
|
|
|
|
func (g *Generator) Close() error {
|
|
f := fmt.Sprintf("%s/data/generator_state.bin", g.HomeDir)
|
|
w, err := os.Create(f)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return g.state.Marshal(w)
|
|
}
|
|
|
|
func encodeUint64(x uint64) []byte {
|
|
var b [8]byte
|
|
b[0] = byte(x)
|
|
b[1] = byte(x >> 8)
|
|
b[2] = byte(x >> 16)
|
|
b[3] = byte(x >> 24)
|
|
b[4] = byte(x >> 32)
|
|
b[5] = byte(x >> 40)
|
|
b[6] = byte(x >> 48)
|
|
b[7] = byte(x >> 56)
|
|
return b[:]
|
|
}
|
|
|
|
const maxStoreKeyGenIterations = 100
|
|
|
|
// StoreKeys deterministically generates a set of unique store keys from seed.
|
|
func StoreKeys(prefix string, seed, count uint64) ([]string, error) {
|
|
r := rand.New(rand.NewPCG(seed, seed>>32))
|
|
keys := make([]string, count)
|
|
seen := make(map[string]struct{})
|
|
|
|
var i, j uint64
|
|
for i < count {
|
|
if j > maxStoreKeyGenIterations {
|
|
return nil, fmt.Errorf("failed to generate %d unique store keys", count)
|
|
}
|
|
sk := fmt.Sprintf("%s_%x", prefix, Bytes(r.Uint64(), 8))
|
|
if _, ok := seen[sk]; ok {
|
|
j++
|
|
continue
|
|
}
|
|
keys[i] = sk
|
|
seen[sk] = struct{}{}
|
|
i++
|
|
j++
|
|
}
|
|
return keys, nil
|
|
}
|
|
|
|
// Bytes generates a byte slice of length length from seed.
|
|
// The byte slice is deterministically generated using the (very fast) xxhash algorithm.
|
|
func Bytes(seed, length uint64) []byte {
|
|
b := make([]byte, length)
|
|
rounds := length / 8
|
|
remainder := length % 8
|
|
var h uint64
|
|
for i := uint64(0); i < rounds; i++ {
|
|
h = xxhash.Sum64(encodeUint64(seed + i))
|
|
for j := uint64(0); j < 8; j++ {
|
|
b[i*8+j] = byte(h >> (8 * j))
|
|
}
|
|
}
|
|
if remainder > 0 {
|
|
h = xxhash.Sum64(encodeUint64(seed + rounds))
|
|
for j := uint64(0); j < remainder; j++ {
|
|
b[rounds*8+j] = byte(h >> (8 * j))
|
|
}
|
|
}
|
|
return b
|
|
}
|