ceremonyclient/pebble/internal/randvar/flag.go

161 lines
3.5 KiB
Go
Raw Normal View History

2024-01-03 07:31:42 +00:00
// Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
// of this source code is governed by a BSD-style license that can be found in
// the LICENSE file.
package randvar
import (
"encoding/binary"
"regexp"
"strconv"
"strings"
"github.com/cockroachdb/errors"
"golang.org/x/exp/rand"
)
var randVarRE = regexp.MustCompile(`^(?:(latest|uniform|zipf):)?(\d+)(?:-(\d+))?$`)
// Flag provides a command line flag interface for specifying static random
// variables.
type Flag struct {
Static
spec string
}
// NewFlag creates a new Flag initialized with the specified spec.
func NewFlag(spec string) *Flag {
f := &Flag{}
if err := f.Set(spec); err != nil {
panic(err)
}
return f
}
func (f *Flag) String() string {
return f.spec
}
// Type implements the Flag.Value interface.
func (f *Flag) Type() string {
return "randvar"
}
// Set implements the Flag.Value interface.
func (f *Flag) Set(spec string) error {
m := randVarRE.FindStringSubmatch(spec)
if m == nil {
return errors.Errorf("invalid random var spec: %s", errors.Safe(spec))
}
min, err := strconv.Atoi(m[2])
if err != nil {
return err
}
max := min
if m[3] != "" {
max, err = strconv.Atoi(m[3])
if err != nil {
return err
}
}
switch strings.ToLower(m[1]) {
case "", "uniform":
f.Static = NewUniform(uint64(min), uint64(max))
case "latest":
f.Static, err = NewSkewedLatest(uint64(min), uint64(max), 0.99)
if err != nil {
return err
}
case "zipf":
var err error
f.Static, err = NewZipf(uint64(min), uint64(max), 0.99)
if err != nil {
return err
}
default:
return errors.Errorf("unknown random var distribution: %s", errors.Safe(m[1]))
}
f.spec = spec
return nil
}
// BytesFlag provides a command line flag interface for specifying random
// bytes. The specification provides for both the length of the random bytes
// and a target compression ratio.
type BytesFlag struct {
sizeFlag Flag
targetCompression float64
spec string
}
// NewBytesFlag creates a new BytesFlag initialized with the specified spec.
func NewBytesFlag(spec string) *BytesFlag {
f := &BytesFlag{}
if err := f.Set(spec); err != nil {
panic(err)
}
return f
}
func (f *BytesFlag) String() string {
return f.spec
}
// Type implements the Flag.Value interface.
func (f *BytesFlag) Type() string {
return "randbytes"
}
// Set implements the Flag.Value interface.
func (f *BytesFlag) Set(spec string) error {
parts := strings.Split(spec, "/")
if len(parts) == 0 || len(parts) > 2 {
return errors.Errorf("invalid randbytes spec: %s", errors.Safe(spec))
}
if err := f.sizeFlag.Set(parts[0]); err != nil {
return err
}
f.targetCompression = 1.0
if len(parts) == 2 {
var err error
f.targetCompression, err = strconv.ParseFloat(parts[1], 64)
if err != nil {
return err
}
}
f.spec = spec
return nil
}
// Bytes returns random bytes. The length of the random bytes comes from the
// internal sizeFlag.
func (f *BytesFlag) Bytes(r *rand.Rand, buf []byte) []byte {
size := int(f.sizeFlag.Uint64(r))
uniqueSize := int(float64(size) / f.targetCompression)
if uniqueSize < 1 {
uniqueSize = 1
}
if cap(buf) < size {
buf = make([]byte, size)
}
data := buf[:size]
offset := 0
for offset+8 <= uniqueSize {
binary.LittleEndian.PutUint64(data[offset:], r.Uint64())
offset += 8
}
word := r.Uint64()
for offset < uniqueSize {
data[offset] = byte(word)
word >>= 8
offset++
}
for offset < size {
data[offset] = data[offset-uniqueSize]
offset++
}
return data
}