<commit_before>package runner

import (
	"fmt"
	"math/rand"
	"net/url"
	"sync"
	"time"

	"github.com/influxdb/influxdb/client"
)

// Timer is struct that can be used to track elaspsed time
type Timer struct {
	start time.Time
	end   time.Time
}

// Start sets a timers `start` field to the current time
func (t *Timer) Start() {
	t.start = time.Now()
}

// Start sets a timers `end` field to the current time
func (t *Timer) Stop() {
	t.end = time.Now()
}

// Elapsed returns the total elapsed time between the `start`
// and `end` fields on a timer.
func (t *Timer) Elapsed() time.Duration {
	return t.end.Sub(t.start)
}

// NewTimer returns a pointer to a `Timer` struct where the
// timers `start` field has been set to `time.Now()`
func NewTimer() *Timer {
	t := &Timer{}
	t.Start()
	return t
}

// Config is a struct that is passed into the `Run()` function.
type Config struct {
	BatchSize     int
	SeriesCount   int
	PointCount    int
	Concurrency   int
	BatchInterval time.Duration
	Database      string
	Address       string
}

// newClient returns a pointer to an InfluxDB client for
// a `Config`'s `Address` field. If an error is encountered
// when creating a new client, the function panics.
func (cfg *Config) newClient() *client.Client {
	u, _ := url.Parse(fmt.Sprintf("http://%s", cfg.Address))
	c, err := client.NewClient(client.Config{URL: *u})
	if err != nil {
		panic(err)
	}
	return c
}

// Run runs the stress test that is specified by a `Config`.
// It returns the total number of points that were during the test,
// an slice of all of the stress tests response times,
// and the times that the test started at and ended as a `Timer`
func Run(cfg *Config) (totalPoints int, responseTimes []int, timer *Timer) {
	timer = NewTimer()
	defer timer.Stop()

	c := cfg.newClient()

	counter := NewConcurrencyLimiter(cfg.Concurrency)

	var mu sync.Mutex
	var wg sync.WaitGroup
	responseTimes = make([]int, 0)

	totalPoints = 0

	batch := &client.BatchPoints{
		Database:         cfg.Database,
		WriteConsistency: "any",
		Time:             time.Now(),
		Precision:        "n",
	}
	for i := 1; i <= cfg.PointCount; i++ {
		for j := 1; j <= cfg.SeriesCount; j++ {
			p := client.Point{
				Measurement: "cpu",
				Tags:        map[string]string{"region": "uswest", "host": fmt.Sprintf("host-%d", j)},
				Fields:      map[string]interface{}{"value": rand.Float64()},
			}
			batch.Points = append(batch.Points, p)
			if len(batch.Points) >= cfg.BatchSize {
				wg.Add(1)
				counter.Increment()
				totalPoints += len(batch.Points)
				go func(b *client.BatchPoints, total int) {
					st := time.Now()
					if _, err := c.Write(*b); err != nil {
						fmt.Println("ERROR: ", err.Error())
					} else {
						mu.Lock()
						responseTimes = append(responseTimes, int(time.Since(st).Nanoseconds()))
						mu.Unlock()
					}
					wg.Done()
					counter.Decrement()
					if total%500000 == 0 {
						fmt.Printf("%d total points. %d in %s\n", total, cfg.BatchSize, time.Since(st))
					}
				}(batch, totalPoints)

				batch = &client.BatchPoints{
					Database:         cfg.Database,
					WriteConsistency: "any",
					Precision:        "n",
					Time:             time.Now(),
				}
			}
		}
	}

	wg.Wait()

	return
}

// ConcurrencyLimiter is a go routine safe struct that can be used to
// ensure that no more than a specifid max number of goroutines are<commit_msg>Change responseTimes to its own type<commit_after>package runner

import (
	"fmt"
	"math/rand"
	"net/url"
	"sync"
	"time"

	"github.com/influxdb/influxdb/client"
)

// Timer is struct that can be used to track elaspsed time
type Timer struct {
	start time.Time
	end   time.Time
}

// Start sets a timers `start` field to the current time
func (t *Timer) Start() {
	t.start = time.Now()
}

// Start sets a timers `end` field to the current time
func (t *Timer) Stop() {
	t.end = time.Now()
}

// Elapsed returns the total elapsed time between the `start`
// and `end` fields on a timer.
func (t *Timer) Elapsed() time.Duration {
	return t.end.Sub(t.start)
}

// NewTimer returns a pointer to a `Timer` struct where the
// timers `start` field has been set to `time.Now()`
func NewTimer() *Timer {
	t := &Timer{}
	t.Start()
	return t
}

// Config is a struct that is passed into the `Run()` function.
type Config struct {
	BatchSize     int
	SeriesCount   int
	PointCount    int
	Concurrency   int
	BatchInterval time.Duration
	Database      string
	Address       string
}

// ResponseTime is a struct that contains `Value`
// `Time` pairing.
type ResponseTime struct {
	Value int
	Time  time.Time
}

// newResponseTime returns a new response time
// with value `v` and time `time.Now()`.
func newResponseTime(v int) ResponseTime {
	r := ResponseTime{Value: v, Time: time.Now()}
	return r
}

type ResponseTimes []ResponseTime

// newClient returns a pointer to an InfluxDB client for
// a `Config`'s `Address` field. If an error is encountered
// when creating a new client, the function panics.
func (cfg *Config) newClient() *client.Client {
	u, _ := url.Parse(fmt.Sprintf("http://%s", cfg.Address))
	c, err := client.NewClient(client.Config{URL: *u})
	if err != nil {
		panic(err)
	}
	return c
}

// Run runs the stress test that is specified by a `Config`.
// It returns the total number of points that were during the test,
// an slice of all of the stress tests response times,
// and the times that the test started at and ended as a `Timer`
func Run(cfg *Config) (totalPoints int, responseTimes ResponseTimes, timer *Timer) {
	timer = NewTimer()
	defer timer.Stop()

	c := cfg.newClient()

	counter := NewConcurrencyLimiter(cfg.Concurrency)

	var mu sync.Mutex
	var wg sync.WaitGroup
	responseTimes = make(ResponseTimes, 0)

	totalPoints = 0

	batch := &client.BatchPoints{
		Database:         cfg.Database,
		WriteConsistency: "any",
		Time:             time.Now(),
		Precision:        "n",
	}
	for i := 1; i <= cfg.PointCount; i++ {
		for j := 1; j <= cfg.SeriesCount; j++ {
			p := client.Point{
				Measurement: "cpu",
				Tags:        map[string]string{"region": "uswest", "host": fmt.Sprintf("host-%d", j)},
				Fields:      map[string]interface{}{"value": rand.Float64()},
			}
			batch.Points = append(batch.Points, p)
			if len(batch.Points) >= cfg.BatchSize {
				wg.Add(1)
				counter.Increment()
				totalPoints += len(batch.Points)
				go func(b *client.BatchPoints, total int) {
					st := time.Now()
					if _, err := c.Write(*b); err != nil {
						fmt.Println("ERROR: ", err.Error())
					} else {
						mu.Lock()
						responseTimes = append(responseTimes, newResponseTime(int(time.Since(st).Nanoseconds())))
						mu.Unlock()
					}
					wg.Done()
					counter.Decrement()
					if total%500000 == 0 {
						fmt.Printf("%d total points. %d in %s\n", total, cfg.BatchSize, time.Since(st))
					}
				}(batch, totalPoints)

				batch = &client.BatchPoints{
					Database:         cfg.Database,
					WriteConsistency: "any",
					Precision:        "n",
					Time:             time.Now(),
				}
			}
		}
	}

	wg.Wait()

	return
}

// ConcurrencyLimiter is a go routine safe struct that can be used to
// ensure that no more than a specifid max number of goroutines are
{'fraction_non_alphanumeric': 0.1325519743267755, 'fraction_numerical': 0.003906795032789173, 'mean_word_length': 2.9602209944751383, 'pattern_counts': {'":': 6, '<': 7, '<?xml version=': 0, '>': 5, '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': '227772', 'n_tokens_mistral': 2652, 'n_tokens_neox': 2349, 'n_words': 786}