forked from TrueCloudLab/xk6-frostfs
117 lines
2.4 KiB
Go
117 lines
2.4 KiB
Go
package datagen
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"math/rand"
|
|
"time"
|
|
|
|
"github.com/dop251/goja"
|
|
"github.com/go-loremipsum/loremipsum"
|
|
"go.k6.io/k6/js/modules"
|
|
)
|
|
|
|
type (
|
|
// Generator stores buffer of random bytes with some tail and returns data slices with
|
|
// an increasing offset so that we receive a different set of bytes after each call without
|
|
// re-generation of the entire buffer from scratch:
|
|
//
|
|
// [<----------size----------><-tail->]
|
|
// [<----------slice0-------->........]
|
|
// [.<----------slice1-------->.......]
|
|
// [..<----------slice2-------->......]
|
|
Generator struct {
|
|
vu modules.VU
|
|
size int
|
|
rand *rand.Rand
|
|
buf []byte
|
|
typ string
|
|
offset int
|
|
}
|
|
|
|
GenPayloadResponse struct {
|
|
Payload goja.ArrayBuffer
|
|
Hash string
|
|
}
|
|
)
|
|
|
|
// TailSize specifies number of extra random bytes in the buffer tail.
|
|
const TailSize = 1024
|
|
|
|
var payloadTypes = []string{
|
|
"text",
|
|
"random",
|
|
"",
|
|
}
|
|
|
|
func NewGenerator(vu modules.VU, size int, typ string) Generator {
|
|
if size <= 0 {
|
|
panic("size should be positive")
|
|
}
|
|
|
|
var found bool
|
|
for i := range payloadTypes {
|
|
if payloadTypes[i] == typ {
|
|
found = true
|
|
break
|
|
}
|
|
}
|
|
if !found {
|
|
vu.InitEnv().Logger.Info("Unknown payload type '%s', random will be used.", typ)
|
|
}
|
|
|
|
r := rand.New(rand.NewSource(time.Now().UnixNano()))
|
|
buf := make([]byte, size+TailSize)
|
|
g := Generator{
|
|
vu: vu,
|
|
size: size,
|
|
rand: r,
|
|
buf: buf,
|
|
typ: typ,
|
|
}
|
|
g.fillBuffer()
|
|
return g
|
|
}
|
|
|
|
func (g *Generator) fillBuffer() {
|
|
switch g.typ {
|
|
case "text":
|
|
li := loremipsum.New()
|
|
b := bytes.NewBuffer(g.buf[:0])
|
|
for b.Len() < g.size+TailSize {
|
|
b.WriteString(li.Paragraph())
|
|
b.WriteRune('\n')
|
|
}
|
|
g.buf = b.Bytes()
|
|
default:
|
|
rand.Read(g.buf) // Per docs, err is always nil here
|
|
}
|
|
}
|
|
|
|
func (g *Generator) GenPayload(calcHash bool) GenPayloadResponse {
|
|
data := g.nextSlice()
|
|
|
|
dataHash := ""
|
|
if calcHash {
|
|
hashBytes := sha256.Sum256(data)
|
|
dataHash = hex.EncodeToString(hashBytes[:])
|
|
}
|
|
|
|
payload := g.vu.Runtime().NewArrayBuffer(data)
|
|
return GenPayloadResponse{Payload: payload, Hash: dataHash}
|
|
}
|
|
|
|
func (g *Generator) nextSlice() []byte {
|
|
if g.offset+g.size >= len(g.buf) {
|
|
g.offset = 0
|
|
g.fillBuffer()
|
|
}
|
|
|
|
result := g.buf[g.offset : g.offset+g.size]
|
|
|
|
// Shift the offset for the next call. If we've used our entire tail, then erase
|
|
// the buffer so that on the next call it is regenerated anew
|
|
g.offset += 1
|
|
return result
|
|
}
|