113 lines
2.3 KiB
Go
113 lines
2.3 KiB
Go
package datagen
|
|
|
|
import (
|
|
"bytes"
|
|
"math/rand"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"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
|
|
|
|
streaming bool
|
|
seed *atomic.Int64
|
|
}
|
|
)
|
|
|
|
// 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, streaming bool) 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)
|
|
}
|
|
g := Generator{
|
|
vu: vu,
|
|
size: size,
|
|
typ: typ,
|
|
}
|
|
|
|
if streaming {
|
|
g.streaming = true
|
|
g.seed = new(atomic.Int64)
|
|
} else {
|
|
g.rand = rand.New(rand.NewSource(time.Now().UnixNano()))
|
|
g.buf = make([]byte, size+TailSize)
|
|
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:
|
|
g.rand.Read(g.buf) // Per docs, err is always nil here
|
|
}
|
|
}
|
|
|
|
func (g *Generator) GenPayload() Payload {
|
|
if g.streaming {
|
|
return NewStreamPayload(g.size, g.seed.Add(1), g.typ)
|
|
}
|
|
|
|
data := g.nextSlice()
|
|
return NewFixedPayload(data)
|
|
}
|
|
|
|
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
|
|
}
|