1
0
mirror of https://github.com/Oxalide/vsphere-influxdb-go.git synced 2023-10-10 13:36:51 +02:00
vsphere-influxdb-go/vendor/github.com/influxdata/influxdb/stress/basic.go
2017-10-25 20:52:40 +00:00

697 lines
14 KiB
Go

package stress
import (
"bytes"
"errors"
"fmt"
"io"
"io/ioutil"
"math/rand"
"net/http"
"sync"
"time"
"github.com/influxdata/influxdb/client/v2"
)
const backoffInterval = time.Duration(500 * time.Millisecond)
// AbstractTag is a struct that abstractly
// defines a tag
type AbstractTag struct {
Key string `toml:"key"`
Value string `toml:"value"`
}
// AbstractTags is a slice of abstract tags
type AbstractTags []AbstractTag
// Template returns a templated string of tags
func (t AbstractTags) Template() string {
var buf bytes.Buffer
for i, tag := range t {
if i == 0 {
buf.Write([]byte(fmt.Sprintf("%v=%v-%%v,", tag.Key, tag.Value)))
} else {
buf.Write([]byte(fmt.Sprintf("%v=%v,", tag.Key, tag.Value)))
}
}
b := buf.Bytes()
b = b[0 : len(b)-1]
return string(b)
}
// AbstractField is a struct that abstractly
// defines a field
type AbstractField struct {
Key string `toml:"key"`
Type string `toml:"type"`
}
// AbstractFields is a slice of abstract fields
type AbstractFields []AbstractField
// Template returns a templated string of fields
func (f AbstractFields) Template() (string, []string) {
var buf bytes.Buffer
a := make([]string, len(f))
for i, field := range f {
buf.Write([]byte(fmt.Sprintf("%v=%%v,", field.Key)))
a[i] = field.Type
}
b := buf.Bytes()
b = b[0 : len(b)-1]
return string(b), a
}
// BasicPointGenerator implements the PointGenerator interface
type BasicPointGenerator struct {
PointCount int `toml:"point_count"`
Tick string `toml:"tick"`
Jitter bool `toml:"jitter"`
Measurement string `toml:"measurement"`
SeriesCount int `toml:"series_count"`
Tags AbstractTags `toml:"tag"`
Fields AbstractFields `toml:"field"`
StartDate string `toml:"start_date"`
Precision string `toml:"precision"`
time time.Time
mu sync.Mutex
}
// typeArr accepts a string array of types and
// returns an array of equal length where each
// element of the array is an instance of the type
// expressed in the string array.
func typeArr(a []string) []interface{} {
i := make([]interface{}, len(a))
for j, ty := range a {
var t string
switch ty {
case "float64":
t = fmt.Sprintf("%v", rand.Intn(1000))
case "int":
t = fmt.Sprintf("%vi", rand.Intn(1000))
case "bool":
b := rand.Intn(2) == 1
t = fmt.Sprintf("%t", b)
default:
t = fmt.Sprintf("%v", rand.Intn(1000))
}
i[j] = t
}
return i
}
func (b *BasicPointGenerator) timestamp(t time.Time) int64 {
var n int64
if b.Precision == "s" {
n = t.Unix()
} else {
n = t.UnixNano()
}
return n
}
// Template returns a function that returns a pointer to a Pnt.
func (b *BasicPointGenerator) Template() func(i int, t time.Time) *Pnt {
ts := b.Tags.Template()
fs, fa := b.Fields.Template()
tmplt := fmt.Sprintf("%v,%v %v %%v", b.Measurement, ts, fs)
return func(i int, t time.Time) *Pnt {
p := &Pnt{}
arr := []interface{}{i}
arr = append(arr, typeArr(fa)...)
arr = append(arr, b.timestamp(t))
str := fmt.Sprintf(tmplt, arr...)
p.Set([]byte(str))
return p
}
}
// Pnt is a struct that implements the Point interface.
type Pnt struct {
line []byte
}
// Set sets the internal state for a Pnt.
func (p *Pnt) Set(b []byte) {
p.line = b
}
// Next generates very simple points very
// efficiently.
// TODO: Take this out
func (p *Pnt) Next(i int, t time.Time) {
p.line = []byte(fmt.Sprintf("a,b=c-%v v=%v", i, i))
}
// Line returns a byte array for a point
// in line protocol format.
func (p Pnt) Line() []byte {
return p.line
}
// Graphite returns a byte array for a point
// in graphite format.
func (p Pnt) Graphite() []byte {
// TODO: Implement
return []byte("")
}
// OpenJSON returns a byte array for a point
// in opentsdb json format
func (p Pnt) OpenJSON() []byte {
// TODO: Implement
return []byte("")
}
// OpenTelnet returns a byte array for a point
// in opentsdb-telnet format
func (p Pnt) OpenTelnet() []byte {
// TODO: Implement
return []byte("")
}
// Generate returns a point channel. Implements the
// Generate method for the PointGenerator interface
func (b *BasicPointGenerator) Generate() (<-chan Point, error) {
// TODO: should be 1.5x batch size
c := make(chan Point, 15000)
tmplt := b.Template()
go func(c chan Point) {
defer close(c)
var start time.Time
var err error
if b.StartDate == "now" {
start = time.Now()
} else {
start, err = time.Parse("2006-Jan-02", b.StartDate)
}
if err != nil {
fmt.Println(err)
return
}
b.mu.Lock()
b.time = start
b.mu.Unlock()
tick, err := time.ParseDuration(b.Tick)
if err != nil {
fmt.Println(err)
return
}
for i := 0; i < b.PointCount; i++ {
b.mu.Lock()
b.time = b.time.Add(tick)
b.mu.Unlock()
for j := 0; j < b.SeriesCount; j++ {
p := tmplt(j, b.time)
c <- *p
}
}
}(c)
return c, nil
}
// Time returns the timestamp for the latest points
// that are being generated. Implements the Time method
// for the PointGenerator interface.
func (b *BasicPointGenerator) Time() time.Time {
defer b.mu.Unlock()
b.mu.Lock()
t := b.time
return t
}
// BasicClient implements the InfluxClient
// interface.
type BasicClient struct {
Enabled bool `toml:"enabled"`
Addresses []string `toml:"addresses"`
Database string `toml:"database"`
RetentionPolicy string `toml:"retention-policy"`
Precision string `toml:"precision"`
BatchSize int `toml:"batch_size"`
BatchInterval string `toml:"batch_interval"`
Concurrency int `toml:"concurrency"`
SSL bool `toml:"ssl"`
Format string `toml:"format"`
addrId int
r chan<- response
interval time.Duration
}
func (c *BasicClient) retry(b []byte, backoff time.Duration) {
bo := backoff + backoffInterval
rs, err := c.send(b)
time.Sleep(c.interval)
c.r <- rs
if !rs.Success() || err != nil {
time.Sleep(bo)
c.retry(b, bo)
}
}
// Batch groups together points
func (c *BasicClient) Batch(ps <-chan Point, r chan<- response) error {
if !c.Enabled {
return nil
}
instanceURLs := make([]string, len(c.Addresses))
for i := 0; i < len(c.Addresses); i++ {
instanceURLs[i] = fmt.Sprintf("http://%v/write?db=%v&rp=%v&precision=%v", c.Addresses[i], c.Database, c.RetentionPolicy, c.Precision)
}
c.Addresses = instanceURLs
c.r = r
var buf bytes.Buffer
var wg sync.WaitGroup
counter := NewConcurrencyLimiter(c.Concurrency)
interval, err := time.ParseDuration(c.BatchInterval)
if err != nil {
return err
}
c.interval = interval
ctr := 0
writeBatch := func(b []byte) {
wg.Add(1)
counter.Increment()
go func(byt []byte) {
c.retry(byt, time.Duration(1))
counter.Decrement()
wg.Done()
}(b)
}
for p := range ps {
b := p.Line()
c.addrId = ctr % len(c.Addresses)
ctr++
buf.Write(b)
buf.Write([]byte("\n"))
if ctr%c.BatchSize == 0 && ctr != 0 {
b := buf.Bytes()
if len(b) == 0 {
continue
}
// Trimming the trailing newline character
b = b[0 : len(b)-1]
writeBatch(b)
var temp bytes.Buffer
buf = temp
}
}
// Write out any remaining points
b := buf.Bytes()
if len(b) > 0 {
writeBatch(b)
}
wg.Wait()
return nil
}
// post sends a post request with a payload of points
func post(url string, datatype string, data io.Reader) (*http.Response, error) {
resp, err := http.Post(url, datatype, data)
if err != nil {
return nil, err
}
defer resp.Body.Close()
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusNoContent && resp.StatusCode != http.StatusOK {
err := errors.New(string(body))
return nil, err
}
return resp, nil
}
// Send calls post and returns a response
func (c *BasicClient) send(b []byte) (response, error) {
t := NewTimer()
resp, err := post(c.Addresses[c.addrId], "", bytes.NewBuffer(b))
t.StopTimer()
if err != nil {
return response{Timer: t}, err
}
r := response{
Resp: resp,
Time: time.Now(),
Timer: t,
}
return r, nil
}
// BasicQuery implements the QueryGenerator interface
type BasicQuery struct {
Template Query `toml:"template"`
QueryCount int `toml:"query_count"`
time time.Time
}
// QueryGenerate returns a Query channel
func (q *BasicQuery) QueryGenerate(now func() time.Time) (<-chan Query, error) {
c := make(chan Query, 0)
go func(chan Query) {
defer close(c)
for i := 0; i < q.QueryCount; i++ {
c <- Query(fmt.Sprintf(string(q.Template), i))
}
}(c)
return c, nil
}
// SetTime sets the internal state of time
func (q *BasicQuery) SetTime(t time.Time) {
q.time = t
return
}
// BasicQueryClient implements the QueryClient interface
type BasicQueryClient struct {
Enabled bool `toml:"enabled"`
Addresses []string `toml:"addresses"`
Database string `toml:"database"`
QueryInterval string `toml:"query_interval"`
Concurrency int `toml:"concurrency"`
clients []client.Client
addrId int
}
// Init initializes the InfluxDB client
func (b *BasicQueryClient) Init() error {
for _, a := range b.Addresses {
cl, err := client.NewHTTPClient(client.HTTPConfig{
Addr: fmt.Sprintf("http://%v", a),
})
if err != nil {
return err
}
b.clients = append(b.clients, cl)
}
return nil
}
// Query runs the query
func (b *BasicQueryClient) Query(cmd Query) (response, error) {
q := client.Query{
Command: string(cmd),
Database: b.Database,
}
t := NewTimer()
_, err := b.clients[b.addrId].Query(q)
t.StopTimer()
if err != nil {
return response{Timer: t}, err
}
// Needs actual response type
r := response{
Time: time.Now(),
Timer: t,
}
return r, nil
}
// Exec listens to the query channel an executes queries as they come in
func (b *BasicQueryClient) Exec(qs <-chan Query, r chan<- response) error {
if !b.Enabled {
return nil
}
var wg sync.WaitGroup
counter := NewConcurrencyLimiter(b.Concurrency)
b.Init()
interval, err := time.ParseDuration(b.QueryInterval)
if err != nil {
return err
}
ctr := 0
for q := range qs {
b.addrId = ctr % len(b.Addresses)
ctr++
wg.Add(1)
counter.Increment()
func(q Query) {
defer wg.Done()
qr, _ := b.Query(q)
r <- qr
time.Sleep(interval)
counter.Decrement()
}(q)
}
wg.Wait()
return nil
}
// resetDB will drop an create a new database on an existing
// InfluxDB instance.
func resetDB(c client.Client, database string) error {
_, err := c.Query(client.Query{
Command: fmt.Sprintf("DROP DATABASE %s", database),
})
if err != nil {
return err
}
_, err = c.Query(client.Query{
Command: fmt.Sprintf("CREATE DATABASE %s", database),
})
if err != nil {
return err
}
return nil
}
// BasicProvisioner implements the Provisioner
// interface.
type BasicProvisioner struct {
Enabled bool `toml:"enabled"`
Address string `toml:"address"`
Database string `toml:"database"`
ResetDatabase bool `toml:"reset_database"`
}
// Provision runs the resetDB function.
func (b *BasicProvisioner) Provision() error {
if !b.Enabled {
return nil
}
cl, err := client.NewHTTPClient(client.HTTPConfig{
Addr: fmt.Sprintf("http://%v", b.Address),
})
if err != nil {
return err
}
if b.ResetDatabase {
resetDB(cl, b.Database)
}
return nil
}
type BroadcastChannel struct {
chs []chan response
wg sync.WaitGroup
fns []func(t *Timer)
}
func NewBroadcastChannel() *BroadcastChannel {
chs := make([]chan response, 0)
return &BroadcastChannel{chs: chs}
}
func (b *BroadcastChannel) Register(fn responseHandler) {
ch := make(chan response, 0)
b.chs = append(b.chs, ch)
f := func(t *Timer) {
go fn(ch, t)
}
b.fns = append(b.fns, f)
}
func (b *BroadcastChannel) Broadcast(r response) {
b.wg.Add(1)
for _, ch := range b.chs {
b.wg.Add(1)
go func(ch chan response) {
ch <- r
b.wg.Done()
}(ch)
}
b.wg.Done()
}
func (b *BroadcastChannel) Close() {
b.wg.Wait()
for _, ch := range b.chs {
close(ch)
// Workaround
time.Sleep(1 * time.Second)
}
}
func (b *BroadcastChannel) Handle(rs <-chan response, t *Timer) {
// Start all of the handlers
for _, fn := range b.fns {
fn(t)
}
for i := range rs {
b.Broadcast(i)
}
b.Close()
}
// BasicWriteHandler handles write responses.
func (b *BasicClient) BasicWriteHandler(rs <-chan response, wt *Timer) {
n := 0
success := 0
fail := 0
s := time.Duration(0)
for t := range rs {
n++
if t.Success() {
success++
} else {
fail++
}
s += t.Timer.Elapsed()
}
if n == 0 {
return
}
fmt.Printf("Total Requests: %v\n", n)
fmt.Printf(" Success: %v\n", success)
fmt.Printf(" Fail: %v\n", fail)
fmt.Printf("Average Response Time: %v\n", s/time.Duration(n))
fmt.Printf("Points Per Second: %v\n\n", int(float64(n)*float64(b.BatchSize)/float64(wt.Elapsed().Seconds())))
}
// BasicReadHandler handles read responses.
func (b *BasicQueryClient) BasicReadHandler(r <-chan response, rt *Timer) {
n := 0
s := time.Duration(0)
for t := range r {
n++
s += t.Timer.Elapsed()
}
if n == 0 {
return
}
fmt.Printf("Total Queries: %v\n", n)
fmt.Printf("Average Query Response Time: %v\n\n", s/time.Duration(n))
}
func (o *outputConfig) HTTPHandler(method string) func(r <-chan response, rt *Timer) {
return func(r <-chan response, rt *Timer) {
c, _ := client.NewHTTPClient(client.HTTPConfig{
Addr: o.addr,
})
bp, _ := client.NewBatchPoints(client.BatchPointsConfig{
Database: o.database,
RetentionPolicy: o.retentionPolicy,
Precision: "ns",
})
for p := range r {
tags := make(map[string]string, len(o.tags))
for k, v := range o.tags {
tags[k] = v
}
tags["method"] = method
fields := map[string]interface{}{
"response_time": float64(p.Timer.Elapsed()),
}
pt, _ := client.NewPoint("performance", tags, fields, p.Time)
bp.AddPoint(pt)
if len(bp.Points())%1000 == 0 && len(bp.Points()) != 0 {
c.Write(bp)
bp, _ = client.NewBatchPoints(client.BatchPointsConfig{
Database: o.database,
RetentionPolicy: o.retentionPolicy,
Precision: "ns",
})
}
}
if len(bp.Points()) != 0 {
c.Write(bp)
}
}
}