package inputs
import (
"bufio"
"fmt"
"io"
"math/rand"
"os"
"sort"
"time"
"github.com/spf13/pflag"
"github.com/timescale/tsbs/cmd/tsbs_generate_data/common"
"github.com/timescale/tsbs/cmd/tsbs_generate_data/devops"
"github.com/timescale/tsbs/cmd/tsbs_generate_data/iot"
"github.com/timescale/tsbs/cmd/tsbs_generate_data/serialize"
)
const (
ErrNoConfig = "no GeneratorConfig provided"
ErrInvalidDataConfig = "invalid config: DataGenerator needs a DataGeneratorConfig"
errLogIntervalZero = "cannot have log interval of 0"
errTotalGroupsZero = "incorrect interleaved groups configuration: total groups = 0"
errInvalidGroupsFmt = "incorrect interleaved groups configuration: id %d >= total groups %d"
errCannotParseTimeFmt = "cannot parse time from string '%s': %v"
)
const defaultLogInterval = 10 * time.Second
type DataGeneratorConfig struct {
BaseConfig
Limit uint64 `mapstructure:"max-data-points"`
InitialScale uint64 `mapstructure:"initial-scale"`
LogInterval time.Duration `mapstructure:"log-interval"`
InterleavedGroupID uint `mapstructure:"interleaved-generation-group-id"`
InterleavedNumGroups uint `mapstructure:"interleaved-generation-groups"`
}
func (c *DataGeneratorConfig) Validate() error {
err := c.BaseConfig.Validate()
if err != nil {
return err
}
if c.InitialScale == 0 {
c.InitialScale = c.BaseConfig.Scale
}
if c.LogInterval == 0 {
return fmt.Errorf(errLogIntervalZero)
}
err = validateGroups(c.InterleavedGroupID, c.InterleavedNumGroups)
return err
}
func (c *DataGeneratorConfig) AddToFlagSet(fs *pflag.FlagSet) {
c.BaseConfig.AddToFlagSet(fs)
fs.Uint64("max-data-points", 0, "Limit the number of data points to generate, 0 = no limit")
fs.Uint64("initial-scale", 0, "Initial scaling variable specific to the use case (e.g., devices in 'devops'). 0 means to use -scale value")
fs.Duration("log-interval", defaultLogInterval, "Duration between data points")
fs.Uint("interleaved-generation-group-id", 0,
"Group (0-indexed) to perform round-robin serialization within. Use this to scale up data generation to multiple processes.")
fs.Uint("interleaved-generation-groups", 1,
"The number of round-robin serialization groups. Use this to scale up data generation to multiple processes.")
}
type DataGenerator struct {
Out io.Writer
config *DataGeneratorConfig
tsStart time.Time
tsEnd time.Time
bufOut *bufio.Writer
}
func (g *DataGenerator) init(config GeneratorConfig) error {
if config == nil {
return fmt.Errorf(ErrNoConfig)
}
switch config.(type) {
case *DataGeneratorConfig:
default:
return fmt.Errorf(ErrInvalidDataConfig)
}
g.config = config.(*DataGeneratorConfig)
err := g.config.Validate()
if err != nil {
return err
}
g.tsStart, err = ParseUTCTime(g.config.TimeStart)
if err != nil {
return fmt.Errorf(errCannotParseTimeFmt, g.config.TimeStart, err)
}
g.tsEnd, err = ParseUTCTime(g.config.TimeEnd)
if err != nil {
return fmt.Errorf(errCannotParseTimeFmt, g.config.TimeEnd, err)
}
if g.Out == nil {
g.Out = os.Stdout
}
g.bufOut, err = getBufferedWriter(g.config.File, g.Out)
if err != nil {
return err
}
return nil
}
func (g *DataGenerator) Generate(config GeneratorConfig) error {
err := g.init(config)
if err != nil {
return err
}
rand.Seed(g.config.Seed)
scfg, err := g.getSimulatorConfig(g.config)
if err != nil {
return err
}
sim := scfg.NewSimulator(g.config.LogInterval, g.config.Limit)
serializer, err := g.getSerializer(sim, g.config.Format)
if err != nil {
return err
}
return g.runSimulator(sim, serializer, g.config)
}
func (g *DataGenerator) runSimulator(sim common.Simulator, serializer serialize.PointSerializer, dgc *DataGeneratorConfig) error {
defer g.bufOut.Flush()
currGroupID := uint(0)
point := serialize.NewPoint()
for !sim.Finished() {
write := sim.Next(point)
if !write {
point.Reset()
continue
}
if currGroupID == dgc.InterleavedGroupID {
err := serializer.Serialize(point, g.bufOut)
if err != nil {
return fmt.Errorf("can not serialize point: %s", err)
}
}
point.Reset()
currGroupID = (currGroupID + 1) % dgc.InterleavedNumGroups
}
return nil
}
func (g *DataGenerator) getSimulatorConfig(dgc *DataGeneratorConfig) (common.SimulatorConfig, error) {
var ret common.SimulatorConfig
var err error
switch dgc.Use {
case useCaseDevops:
ret = &devops.DevopsSimulatorConfig{
Start: g.tsStart,
End: g.tsEnd,
InitHostCount: dgc.InitialScale,
HostCount: dgc.Scale,
HostConstructor: devops.NewHost,
}
case useCaseIoT:
ret = &iot.SimulatorConfig{
Start: g.tsStart,
End: g.tsEnd,
InitGeneratorScale: dgc.InitialScale,
GeneratorScale: dgc.Scale,
GeneratorConstructor: iot.NewTruck,
}
case useCaseCPUOnly:
ret = &devops.CPUOnlySimulatorConfig{
Start: g.tsStart,
End: g.tsEnd,
InitHostCount: dgc.InitialScale,
HostCount: dgc.Scale,
HostConstructor: devops.NewHostCPUOnly,
}
case useCaseCPUSingle:
ret = &devops.CPUOnlySimulatorConfig{
Start: g.tsStart,
End: g.tsEnd,
InitHostCount: dgc.InitialScale,
HostCount: dgc.Scale,
HostConstructor: devops.NewHostCPUSingle,
}
default:
err = fmt.Errorf("unknown use case: '%s'", dgc.Use)
}
return ret, err
}
func (g *DataGenerator) getSerializer(sim common.Simulator, format string) (serialize.PointSerializer, error) {
var ret serialize.PointSerializer
var err error
switch format {
case FormatCassandra:
ret = &serialize.CassandraSerializer{}
case FormatInflux:
ret = &serialize.InfluxSerializer{}
case FormatMongo:
ret = &serialize.MongoSerializer{}
case FormatSiriDB:
ret = &serialize.SiriDBSerializer{}
case FormatAkumuli:
ret = serialize.NewAkumuliSerializer()
case FormatCrateDB:
g.writeHeader(sim)
ret = &serialize.CrateDBSerializer{}
case FormatClickhouse:
fallthrough
case FormatTimescaleDB:
g.writeHeader(sim)
ret = &serialize.TimescaleDBSerializer{}
default:
err = fmt.Errorf(errUnknownFormatFmt, format)
}
return ret, err
}
func (g *DataGenerator) writeHeader(sim common.Simulator) {
g.bufOut.WriteString("tags")
types := sim.TagTypes()
for i, key := range sim.TagKeys() {
g.bufOut.WriteString(",")
g.bufOut.Write(key)
g.bufOut.WriteString(" ")
g.bufOut.WriteString(types[i].String())
}
g.bufOut.WriteString("\n")
keys := make([]string, 0)
fields := sim.Fields()
for k := range fields {
keys = append(keys, k)
}
sort.Strings(keys)
for _, measurementName := range keys {
g.bufOut.WriteString(measurementName)
for _, field := range fields[measurementName] {
g.bufOut.WriteString(",")
g.bufOut.Write(field)
}
g.bufOut.WriteString("\n")
}
g.bufOut.WriteString("\n")
}