<commit_before>package runner

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"time"

	"github.com/grafana/metrictank/cmd/mt-simulate-lock-congestion/metricname"
	"github.com/grafana/metrictank/cmd/mt-simulate-lock-congestion/query"
	"github.com/grafana/metrictank/conf"
	"github.com/grafana/metrictank/input"
	"github.com/grafana/metrictank/mdata"

	"github.com/grafana/metrictank/idx"
	"github.com/grafana/metrictank/idx/memory"
	"github.com/raintank/schema"
	"github.com/raintank/schema/msg"
	"golang.org/x/sync/errgroup"
)

type addedMetricData struct {
	md        *schema.MetricData
	partition int32
}

// TestRun represents one test run, including all the structures used for the test
type TestRun struct {
	index                 idx.MetricIndex
	handler               input.Handler
	metricNameGenerator   metricname.NameGenerator
	queryPatternGenerator query.QueryGenerator
	addedMdChan           chan addedMetricData
	addedMdLock           sync.RWMutex
	addedMds              []addedMetricData
	stats                 runStats
	addsPerSec            uint32 // how many adds per second we want to execute. 0 means unlimited, as many as possible
	addThreads            uint32 // number of concurrent add threads
	addSampleFactor       uint32 // how often to print info about a completed add
	addDelay              uint32 // once the benchmark starts the querying threads start, if addDelay is >0 then the adding into the index is delayed by the given number of seconds
	updatesPerAdd         uint32 // how many metric updates should be done between each metric add, they also get called by the add threads
	initialIndexSize      uint32 // prepopulate the index with the defined number of entries before starting the actual test run
	queriesPerSec         uint32 // how many queries per second we want to execute. 0 means unlimited, as many as possible
	querySampleFactor     uint32 // how often to print info about a completed query
	totalQueryCount       uint32 // the total number of queries we will start
	runDuration           time.Duration
}

const orgID = 1

// NewTestRun Instantiates a new test run
func NewTestRun(nameGenerator metricname.NameGenerator, queryGenerator query.QueryGenerator, addDelay, addsPerSec, addThreads, addSampleFactor, updatesPerAdd, initialIndexSize, queriesPerSec, querySampleFactor uint32, runDuration time.Duration) *TestRun {
	totalQueryCount := (queriesPerSec + addDelay) * uint32(runDuration.Seconds())

	index := memory.New()
	// initializing with a `nil` store, that's a bit risky but good enough for the moment
	mdata.Schemas = conf.NewSchemas(nil)
	metrics := mdata.NewAggMetrics(nil, nil, false, 3600, 7200, 3600)
	handler := input.NewDefaultHandler(metrics, index, "lock-congestion-simulator")

	runner := TestRun{
		index:                 index,
		handler:               handler,
		stats:                 newRunStats(totalQueryCount),
		totalQueryCount:       totalQueryCount,
		metricNameGenerator:   nameGenerator,
		queryPatternGenerator: queryGenerator,
		addedMdChan:           make(chan addedMetricData, 100),
		addsPerSec:            addsPerSec,
		addThreads:            addThreads,
		addSampleFactor:       addSampleFactor,
		addDelay:              addDelay,
		updatesPerAdd:         updatesPerAdd,
		initialIndexSize:      initialIndexSize,
		queriesPerSec:         queriesPerSec,
		querySampleFactor:     querySampleFactor,
		runDuration:           runDuration,
	}

	return &runner
}

// Run executes the run
func (t *TestRun) Run() {
	mainCtx, cancel := context.WithCancel(context.Background())

	// wait group to wait until all routines have been started
	startedWg := sync.WaitGroup{}
	workerThreads, workerCtx := errgroup.WithContext(mainCtx)

	startedWg.Add(1)
	workerThreads.Go(t.addedMdTracker(workerCtx, &startedWg))

	log.Printf("Starting the metric name generator")
	t.metricNameGenerator.Start(mainCtx, t.addThreads)

	log.Printf("Starting the query generator")
	t.queryPatternGenerator.Start()

	log.Printf("Prepopulating the index with %d entries", t.initialIndexSize)
	t.prepopulateIndex()

	startedWg.Add(1)
	workerThreads.Go(t.queryRoutine(workerCtx, &startedWg))

	<-time.After(time.Second * time.Duration(t.addDelay))

	startedWg.Add(int(t.addThreads))
	for i := uint32(0); i < t.addThreads; i++ {
		workerThreads.Go(t.addRoutine(workerCtx, &startedWg, i))
	}
	startedWg.Wait()
	log.Printf("Benchmark has started")

	// as soon as all routines have started, we start the timer
	<-time.After(t.runDuration)
	cancel()
	workerThreads.Wait()
}

// PrintStats writes all the statistics in human readable format into stdout
func (t *TestRun) PrintStats() {
	t.stats.Print(uint32(t.runDuration.Seconds()))
}

func (t *TestRun) addedMdTracker(ctx context.Context, startedWg *sync.WaitGroup) func() error {
	buffer := make([]addedMetricData, 0, 1000)
	return func() error {
		startedWg.Done()
		for {
			select {
			case md := <-t.addedMdChan:
				if len(buffer) < cap(buffer) {
					buffer = append(buffer, md)
				} else {
					t.addedMdLock.Lock()
					for i := range buffer {
						t.addedMds = append(t.addedMds, buffer[i])
					}
					t.addedMdLock.Unlock()
					buffer = buffer[:0]
				}
			case <-ctx.Done():
				fmt.Println("addedMdTracker is shutting down")
				return nil
			}
		}
	}
}

func (t *TestRun) getAddedMds(count int) []addedMetricData {
	res := make([]addedMetricData, count)
	t.addedMdLock.RLock()
	defer t.addedMdLock.RUnlock()

	addedMdsLen := len(t.addedMds)
	if count > addedMdsLen || addedMdsLen == 0 {<commit_msg>waste less memory<commit_after>package runner

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"time"

	"github.com/grafana/metrictank/cmd/mt-simulate-lock-congestion/metricname"
	"github.com/grafana/metrictank/cmd/mt-simulate-lock-congestion/query"
	"github.com/grafana/metrictank/conf"
	"github.com/grafana/metrictank/input"
	"github.com/grafana/metrictank/mdata"

	"github.com/grafana/metrictank/idx"
	"github.com/grafana/metrictank/idx/memory"
	"github.com/raintank/schema"
	"github.com/raintank/schema/msg"
	"golang.org/x/sync/errgroup"
)

type addedMetricData struct {
	md        *schema.MetricData
	partition int32
}

// TestRun represents one test run, including all the structures used for the test
type TestRun struct {
	index                 idx.MetricIndex
	handler               input.Handler
	metricNameGenerator   metricname.NameGenerator
	queryPatternGenerator query.QueryGenerator
	addedMdChan           chan addedMetricData
	addedMdLock           sync.RWMutex
	addedMds              []addedMetricData
	stats                 runStats
	addsPerSec            uint32 // how many adds per second we want to execute. 0 means unlimited, as many as possible
	addThreads            uint32 // number of concurrent add threads
	addSampleFactor       uint32 // how often to print info about a completed add
	addDelay              uint32 // once the benchmark starts the querying threads start, if addDelay is >0 then the adding into the index is delayed by the given number of seconds
	updatesPerAdd         uint32 // how many metric updates should be done between each metric add, they also get called by the add threads
	initialIndexSize      uint32 // prepopulate the index with the defined number of entries before starting the actual test run
	queriesPerSec         uint32 // how many queries per second we want to execute. 0 means unlimited, as many as possible
	querySampleFactor     uint32 // how often to print info about a completed query
	totalQueryCount       uint32 // the total number of queries we will start
	runDuration           time.Duration
}

const orgID = 1

// NewTestRun Instantiates a new test run
func NewTestRun(nameGenerator metricname.NameGenerator, queryGenerator query.QueryGenerator, addDelay, addsPerSec, addThreads, addSampleFactor, updatesPerAdd, initialIndexSize, queriesPerSec, querySampleFactor uint32, runDuration time.Duration) *TestRun {
	totalQueryCount := (queriesPerSec + addDelay) * uint32(runDuration.Seconds())

	index := memory.New()
	// initializing with a `nil` store, that's a bit risky but good enough for the moment
	mdata.Schemas = conf.NewSchemas(nil)
	metrics := mdata.NewAggMetrics(nil, nil, false, 3600, 7200, 3600)
	handler := input.NewDefaultHandler(metrics, index, "lock-congestion-simulator")

	runner := TestRun{
		index:                 index,
		handler:               handler,
		stats:                 newRunStats(totalQueryCount),
		totalQueryCount:       totalQueryCount,
		metricNameGenerator:   nameGenerator,
		queryPatternGenerator: queryGenerator,
		addedMdChan:           make(chan addedMetricData, 100),
		addsPerSec:            addsPerSec,
		addThreads:            addThreads,
		addSampleFactor:       addSampleFactor,
		addDelay:              addDelay,
		updatesPerAdd:         updatesPerAdd,
		initialIndexSize:      initialIndexSize,
		queriesPerSec:         queriesPerSec,
		querySampleFactor:     querySampleFactor,
		runDuration:           runDuration,
	}

	return &runner
}

// Run executes the run
func (t *TestRun) Run() {
	mainCtx, cancel := context.WithCancel(context.Background())

	// wait group to wait until all routines have been started
	startedWg := sync.WaitGroup{}
	workerThreads, workerCtx := errgroup.WithContext(mainCtx)

	startedWg.Add(1)
	workerThreads.Go(t.addedMdTracker(workerCtx, &startedWg))

	log.Printf("Starting the metric name generator")
	t.metricNameGenerator.Start(mainCtx, t.addThreads)

	log.Printf("Starting the query generator")
	t.queryPatternGenerator.Start()

	log.Printf("Prepopulating the index with %d entries", t.initialIndexSize)
	t.prepopulateIndex()

	startedWg.Add(1)
	workerThreads.Go(t.queryRoutine(workerCtx, &startedWg))

	<-time.After(time.Second * time.Duration(t.addDelay))

	startedWg.Add(int(t.addThreads))
	for i := uint32(0); i < t.addThreads; i++ {
		workerThreads.Go(t.addRoutine(workerCtx, &startedWg, i))
	}
	startedWg.Wait()
	log.Printf("Benchmark has started")

	// as soon as all routines have started, we start the timer
	<-time.After(t.runDuration)
	cancel()
	workerThreads.Wait()
}

// PrintStats writes all the statistics in human readable format into stdout
func (t *TestRun) PrintStats() {
	t.stats.Print(uint32(t.runDuration.Seconds()))
}

func (t *TestRun) addedMdTracker(ctx context.Context, startedWg *sync.WaitGroup) func() error {
	buffer := make([]addedMetricData, 0, 1000)
	trackUpTo := 1000000 // how many added MDs to track at max
	return func() error {
		startedWg.Done()
		for {
			select {
			case md := <-t.addedMdChan:
				if len(buffer) < cap(buffer) {
					buffer = append(buffer, md)
				} else {
					t.addedMdLock.Lock()
					if len(t.addedMds) < trackUpTo {
						for i := range buffer {
							t.addedMds = append(t.addedMds, buffer[i])
						}
					} else {
						startingAt := rand.Intn(trackUpTo)
						for i := 0; i < len(buffer); i++ {
							t.addedMds[(startingAt+i)%trackUpTo] = buffer[i]
						}
					}
					t.addedMdLock.Unlock()
					buffer = buffer[:0]
				}
			case <-ctx.Done():
				fmt.Println("addedMdTracker is shutting down")
				return nil
			}
		}
	}
}

func (t *TestRun) getAddedMds(count int) []addedMetricData {
	res := make([]addedMetricData, count)
	t.addedMdLock.RLock()
	defer t.addedMdLock.RUnlock()

	addedMdsLen := len(t.addedMds)
	if count > addedMdsLen || addedMdsLen == 0 {
{'fraction_non_alphanumeric': 0.09598727240586884, 'fraction_numerical': 0.010783100583348065, 'mean_word_length': 3.4741004349545275, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1884315', 'n_tokens_mistral': 3859, 'n_tokens_neox': 3522, 'n_words': 956}