metricsbench.go

package main

import (
	"fmt"
	"flag"
	"strings"
	"os"
	"strconv"
	"log"
	"math/rand"
)

type Metric struct {
	Measurement string
	Tags        map[string]string
	ValueKeys   []string
}

func (m Metric) String() string {
	tags := make([]string, 1, len(m.Tags) + 1)
	tags[0] = m.Measurement
	for tag, val := range m.Tags {
		tags = append(tags, fmt.Sprintf("%v=%v", tag, val))
	}
	return strings.Join(tags, ",")
}

type Datapoint struct {
	Metric    *Metric
	Values []float64
	Timestamp int
}

func (d Datapoint) String() (s string) {
	if d.Metric != nil && len(d.Values) != 0 {
		keys := make([]string, 0, len(d.Metric.ValueKeys))
		for i, key := range d.Metric.ValueKeys {
			keys = append(keys, fmt.Sprintf("%v=%v", key, d.Values[i]))
		}
		s = fmt.Sprintf("%s %s", d.Metric, strings.Join(keys, ","))
	}
	return
}

func TagClosureFabric() func(int) func() (string, string) {
	var tagId int

	return func(length int) func() (string, string) {
		var i int

		name := fmt.Sprintf("tag%d", tagId)
		tagId++

		return func() (string, string) {
			value := fmt.Sprintf("%d", i)
			i++
			if i > length {
				i = 0
			}

			return name, value
		}
	}
}

func MakeTags(amount int, fields []int) (r []map[string]string) {
	fields_count := len(fields)
	closures := make([]func() (string, string), 0, fields_count)
	fabric := TagClosureFabric()

	for _, length := range fields {
		closures = append(closures, fabric(length))
	}

	r = make([]map[string]string, 0, amount)

	for i := 0; i < amount; i++ {
		res := make(map[string]string, fields_count + 1)
		res["host"] = fmt.Sprintf("host-%d", i)
		for _, f := range closures {
			tag, val := f()
			res[tag] = val
		}
		r = append(r, res)
	}
	return
}

func MakeMetrics(m *map[string][]string, t *[]map[string]string) (*[]Metric) {
	r := make([]Metric, 0, len(*m) * len(*t))
	for _, t := range *t {
		for m, v := range *m {
			metric := Metric{Measurement: m, Tags: t, ValueKeys: v}
			r = append(r,  metric)
		}
	}
	return &r
}

func intArg(s string, errmsg string) (r int) {
	r, err := strconv.Atoi(s)
	if err != nil {
        	log.Println(errmsg);
        	os.Exit(1);
	}
	return
}

func MetricFeeder(metrics *[]Metric, passes int, speed int, output chan *Metric) {
	block := len(*metrics) / speed

	for s := 0; s < speed; s++ {
		stop := (s + 1) * block
		for p := 0; p < passes; p++ {
			for i := s * block; i < stop; i++ {
				output <- &((*metrics)[i])
			}
		}
	}
	close(output)
}

func DataPointWriter(batchSize int, inputMetrics chan *Metric) {
	output := make([]string, batchSize)
	i := 0
	for m := range inputMetrics {
		l := len(m.ValueKeys)
		v := make([]float64, l)
		for j := 0; j < l; j++ {
			v[j] = rand.Float64()
		}
		d := Datapoint{Metric: m, Values: v}
		output[i] = d.String()
		i++

		if i == batchSize {
			outstr := strings.Join(output, "\n")
			fmt.Printf("%d /write?db=test\r\n", len(outstr))
			fmt.Println(outstr)
			output = make([]string, batchSize)
			i = 0
		}
	}
}

func main() {
	var Measurements = map[string][]string{
		"memory_used": {"value"},
		"memory_free": {"value"},
		"memory_cached": {"value"},
		"net_tx": {"value"},
		"net_rx": {"value"},
		"net_recv": {"value"},
		"net_send": {"value"},
		"disk_read": {"value"},
		"disk_write": {"value"},
		"cpu_user": {"value"},
		"cpu_system": {"value"},
		"cpu_iowait": {"value"},
	}

	flag.Parse()
	args := flag.Args()

    	if len(args) < 4 {
		log.Println("Metrics amount is required.")
        	os.Exit(1);
	}

	num_tags := intArg(args[0], "Metrics amount should be a number.")
	passes := intArg(args[1], "Passes amount should be a number.")
	speed := intArg(args[2], "Speed should be a number.")
	batch := intArg(args[3], "Batch size should be a number.")

	fields := make([]int, 0, len(args[4:]))
	var c int = 1
	for _, i := range args[4:] {
		x := intArg(i, "Tag factor should be a number.")
		c = c*x
		fields = append(fields, x)
	}

	tags := MakeTags(num_tags / len(Measurements), fields)
	metrics := MakeMetrics(&Measurements, &tags)
	log.Printf("Generated %d metrics, passes %d, speed %d, batch %d", len(*metrics), passes, speed, batch)
	log.Printf("Total cardinality %d", c * len(*metrics) * len(Measurements))

	metricChannel := make(chan *Metric, 100)

	go MetricFeeder(metrics, passes, speed, metricChannel)
	DataPointWriter(batch, metricChannel)
}