qb/examples/loadtest/loadtest.go

package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"time"

	"gordenko.dev/dima/qb"
	"gordenko.dev/dima/qb/client"
	"gordenko.dev/dima/qb/proto"
)

// METRICS INFO

type MetricInfo struct {
	MetricID   uint32        `json:"metricID"`
	MetricType qb.MetricType `json:"metricType"`
	FracDigits int           `json:"fracDigits"`
	Since      int64         `json:"since"`
	Until      int64         `json:"until"`
	Qty        int           `json:"qty"`
}

func readMetricInfo(fileName string) (list []MetricInfo, err error) {
	buf, err := os.ReadFile(fileName)
	if err != nil {
		return
	}

	err = json.Unmarshal(buf, &list)
	return
}

// RANDOM QUERY GENERATOR

type QueryRecipe struct {
	MetricID  uint32
	MetricIDs []uint32
	Method    int
	RangeCode int
	Since     uint32
	Until     uint32
	GroupBy   qb.GroupBy
}

type RandomQueryGenerator struct {
	metrics                      []MetricInfo
	groupByOptions               []qb.GroupBy
	listCurrentValuesProbability int
	listPeriodsProbability       int
	timeRangeProbDistribution    []int
}

type RandomQueryGeneratorOptions struct {
	Metrics                      []MetricInfo
	ListCurrentValuesProbability int
	ListPeriodsProbability       int
	TimeRangeProbabilities       []int
}

func NewRandomQueryGenerator(opt RandomQueryGeneratorOptions) *RandomQueryGenerator {
	if opt.ListCurrentValuesProbability >= 100 {
		panic(fmt.Sprintf("wrong ListCurrentValuesProbability: %d", opt.ListCurrentValuesProbability))
	}

	if opt.ListPeriodsProbability >= 100 {
		panic(fmt.Sprintf("wrong ListPeriodsProbability: %d", opt.ListPeriodsProbability))
	}
	// check total time range propability
	var totalTimeRangeProbability int
	for _, p := range opt.TimeRangeProbabilities {
		totalTimeRangeProbability += p
	}
	if totalTimeRangeProbability != 100 {
		panic(fmt.Sprintf("total time range probabilities != 100: %d", totalTimeRangeProbability))
	}

	// create time range probability distribution
	timeRangeProbDistribution := make([]int, len(opt.TimeRangeProbabilities))
	timeRangeProbDistribution[0] = opt.TimeRangeProbabilities[0]
	for i := 1; i < len(opt.TimeRangeProbabilities); i++ {
		timeRangeProbDistribution[i] = timeRangeProbDistribution[i-1] + opt.TimeRangeProbabilities[i]
	}

	return &RandomQueryGenerator{
		metrics: opt.Metrics,
		groupByOptions: []qb.GroupBy{
			qb.GroupByHour,
			qb.GroupByDay,
			qb.GroupByMonth,
		},
		listCurrentValuesProbability: opt.ListCurrentValuesProbability,
		listPeriodsProbability:       opt.ListPeriodsProbability,
		timeRangeProbDistribution:    timeRangeProbDistribution,
	}
}

func (s *RandomQueryGenerator) GetQueryRecipe() QueryRecipe {
	metric := s.getRandomMetric()

	num := rand.Intn(100)
	if num < s.listCurrentValuesProbability {
		qty := 5 + rand.Intn(100) // від 5 до 105
		return QueryRecipe{
			MetricIDs: s.listRandomUniqueMetricIDs(qty),
			Method:    listCurrentValues,
		}
	} else {
		if metric.MetricType == qb.Cumulative {
			num = rand.Intn(100)
			if num < s.listPeriodsProbability {
				groupBy := s.groupByOptions[rand.Intn(len(s.groupByOptions))]

				var (
					minDays = 1
					maxDays = 7
				)

				if groupBy == qb.GroupByDay {
					minDays = 1
					maxDays = 30
				} else if groupBy == qb.GroupByMonth {
					minDays = 1
					maxDays = 30
				}

				rangeCode, since, until := s.getRandomTimeRange(
					metric.Since, metric.Until, minDays, maxDays)

				return QueryRecipe{
					MetricID:  metric.MetricID,
					Method:    listCumulativePeriods,
					RangeCode: rangeCode,
					Since:     uint32(since),
					Until:     uint32(until),
					GroupBy:   groupBy,
				}
			} else {
				var (
					minDays = 1
					maxDays = 3
				)

				rangeCode, since, until := s.getRandomTimeRange(
					metric.Since, metric.Until, minDays, maxDays)

				return QueryRecipe{
					MetricID:  metric.MetricID,
					Method:    listCumulativeMeasures,
					RangeCode: rangeCode,
					Since:     uint32(since),
					Until:     uint32(until),
				}
			}
		} else {
			num = rand.Intn(100)
			if num < s.listPeriodsProbability {
				groupBy := s.groupByOptions[rand.Intn(len(s.groupByOptions))]

				var (
					minDays = 1
					maxDays = 7
				)

				if groupBy == qb.GroupByDay {
					minDays = 1
					maxDays = 30
				} else if groupBy == qb.GroupByMonth {
					minDays = 1
					maxDays = 30
				}

				rangeCode, since, until := s.getRandomTimeRange(
					metric.Since, metric.Until, minDays, maxDays)

				return QueryRecipe{
					MetricID:  metric.MetricID,
					Method:    listInstantPeriods,
					RangeCode: rangeCode,
					Since:     uint32(since),
					Until:     uint32(until),
					GroupBy:   groupBy,
				}
			} else {
				var (
					minDays = 1
					maxDays = 3
				)

				rangeCode, since, until := s.getRandomTimeRange(
					metric.Since, metric.Until, minDays, maxDays)

				return QueryRecipe{
					MetricID:  metric.MetricID,
					Method:    listInstantMeasures,
					RangeCode: rangeCode,
					Since:     uint32(since),
					Until:     uint32(until),
				}
			}
		}
	}

}

// Генерує випадковий набір унікальних metricID з [1, 100]
func (s *RandomQueryGenerator) listRandomUniqueMetricIDs(count int) []uint32 {
	// переставляю індекси у випадковому порядку
	indexes := rand.Perm(len(s.metrics))
	// копіюю metricID із перших випадкових індексів
	metricIDs := make([]uint32, count)
	for i := range count {
		metricIDs[i] = s.metrics[indexes[i]].MetricID
	}
	return metricIDs
}

const (
	secondsPerDay   = 86400
	dayRange        = 0
	weekRange       = 1
	monthRange      = 2
	randomTimeRange = 3
)

// Випадковий часовий діапазон
func (s *RandomQueryGenerator) getRandomTimeRange(start, end int64, minDays, maxDays int) (int, int64, int64) {
	var (
		since     int64
		until     int64
		num       = rand.Intn(100)
		rangeCode int
		threshold int
	)
	for rangeCode, threshold = range s.timeRangeProbDistribution {
		if num < threshold {
			break
		}
	}

	switch rangeCode {
	case dayRange:
		since = end - secondsPerDay
		until = end

	case weekRange:
		since = end - 7*secondsPerDay
		until = end

	case monthRange:
		since = end - 30*secondsPerDay
		until = end

	case randomTimeRange:
		if start == end {
			return rangeCode, start, end
		}
		// Випадковий момент часу для since
		since = start + rand.Int63n(end-start)
		// Випадкова тривалість у днях (але не виходити за межу end)
		durationInDays := minDays + rand.Intn(maxDays-minDays)

		until = since + int64(durationInDays)*secondsPerDay
		if until > end {
			until = end
		}
	}
	return rangeCode, since, until
}

func (s *RandomQueryGenerator) getRandomMetric() MetricInfo {
	return s.metrics[rand.Intn(len(s.metrics))]
}

// EXECUTE QUERY

func execQuery(conn *client.Connection, queryGenerator *RandomQueryGenerator, stat *WorkerStat) (err error) {
	recipe := queryGenerator.GetQueryRecipe()

	var elapsedTime time.Duration

	switch recipe.Method {
	case listCurrentValues:
		t1 := time.Now()
		_, err := conn.ListCurrentValues(recipe.MetricIDs)
		elapsedTime = time.Since(t1)
		stat.ElapsedTime += elapsedTime
		stat.Queries++
		stat.ElapsedTimeByMethods[recipe.Method] += elapsedTime
		stat.MethodCalls[recipe.Method]++
		if err != nil {
			return fmt.Errorf("ListCurrentValues: %s", err)
		}

	case listInstantMeasures:
		t1 := time.Now()
		_, err := conn.ListInstantMeasures(proto.ListInstantMeasuresReq{
			MetricID: recipe.MetricID,
			Since:    recipe.Since,
			Until:    recipe.Until,
		})
		elapsedTime = time.Since(t1)
		stat.ElapsedTime += elapsedTime
		stat.Queries++
		stat.ElapsedTimeByMethods[recipe.Method] += elapsedTime
		stat.MethodCalls[recipe.Method]++
		stat.ElapsedTimeByTimeRanges[recipe.RangeCode] += elapsedTime
		stat.TimeRangeCalls[recipe.RangeCode]++
		if err != nil {
			return fmt.Errorf("ListInstantMeasures(%d): %s",
				recipe.MetricID, err)
		}

	case listCumulativeMeasures:
		t1 := time.Now()
		_, err := conn.ListCumulativeMeasures(proto.ListCumulativeMeasuresReq{
			MetricID: recipe.MetricID,
			Since:    recipe.Since,
			Until:    recipe.Until,
		})
		elapsedTime = time.Since(t1)
		stat.ElapsedTime += elapsedTime
		stat.Queries++
		stat.ElapsedTimeByMethods[recipe.Method] += elapsedTime
		stat.MethodCalls[recipe.Method]++
		stat.ElapsedTimeByTimeRanges[recipe.RangeCode] += elapsedTime
		stat.TimeRangeCalls[recipe.RangeCode]++
		if err != nil {
			return fmt.Errorf("ListCumulativeMeasures(%d): %s",
				recipe.MetricID, err)
		}

	case listInstantPeriods:
		since := time.Unix(int64(recipe.Since), 0)
		until := time.Unix(int64(recipe.Until), 0)
		//
		t1 := time.Now()
		_, err := conn.ListInstantPeriods(proto.ListInstantPeriodsReq{
			MetricID: recipe.MetricID,
			Since: proto.TimeBound{
				Year:  since.Year(),
				Month: since.Month(),
				Day:   since.Day(),
			},
			Until: proto.TimeBound{
				Year:  until.Year(),
				Month: until.Month(),
				Day:   until.Day(),
			},
			GroupBy:        recipe.GroupBy,
			AggregateFuncs: qb.AggregateMin | qb.AggregateMax | qb.AggregateAvg,
		})
		elapsedTime = time.Since(t1)
		stat.ElapsedTime += elapsedTime
		stat.Queries++
		stat.ElapsedTimeByMethods[recipe.Method] += elapsedTime
		stat.MethodCalls[recipe.Method]++
		stat.ElapsedTimeByTimeRanges[recipe.RangeCode] += elapsedTime
		stat.TimeRangeCalls[recipe.RangeCode]++
		if err != nil {
			return fmt.Errorf("ListInstantPeriods(%d): %s",
				recipe.MetricID, err)
		}

	case listCumulativePeriods:
		since := time.Unix(int64(recipe.Since), 0)
		until := time.Unix(int64(recipe.Until), 0)
		//
		t1 := time.Now()
		_, err := conn.ListCumulativePeriods(proto.ListCumulativePeriodsReq{
			MetricID: recipe.MetricID,
			Since: proto.TimeBound{
				Year:  since.Year(),
				Month: since.Month(),
				Day:   since.Day(),
			},
			Until: proto.TimeBound{
				Year:  until.Year(),
				Month: until.Month(),
				Day:   until.Day(),
			},
			GroupBy: recipe.GroupBy,
		})
		elapsedTime = time.Since(t1)
		stat.ElapsedTime += elapsedTime
		stat.Queries++
		stat.ElapsedTimeByMethods[recipe.Method] += elapsedTime
		stat.MethodCalls[recipe.Method]++
		stat.ElapsedTimeByTimeRanges[recipe.RangeCode] += elapsedTime
		stat.TimeRangeCalls[recipe.RangeCode]++
		if err != nil {
			return fmt.Errorf("ListCumulativePeriods(%d): %s",
				recipe.MetricID, err)
		}
	}
	return
}