timingbench/bench.go

package timingbench

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"time"

	"github.com/schollz/progressbar/v3"
	"gonum.org/v1/gonum/stat"
)

// TimingResult holds the timing results for a function execution.
type TimingResult struct {
	Min        time.Duration
	Max        time.Duration
	Mean       time.Duration
	Median     time.Duration
	StdDev     time.Duration
	StartTime  time.Time
	EndTime    time.Time
	Duration   time.Duration
	Iterations int
}

// String returns a formatted string representation of the TimingResult.
func (r TimingResult) String() string {
	return fmt.Sprintf(
		"Timing Results:\n"+
			"Start Time: %v\n"+
			"End Time: %v\n"+
			"Total Duration: %v\n"+
			"Iterations: %d\n"+
			"Min Duration: %v\n"+
			"Max Duration: %v\n"+
			"Mean Duration: %v\n"+
			"Median Duration: %v\n"+
			"Standard Deviation: %v\n",
		r.StartTime.Format(time.RFC3339),
		r.EndTime.Format(time.RFC3339),
		r.Duration,
		r.Iterations,
		r.Min,
		r.Max,
		r.Mean,
		r.Median,
		r.StdDev,
	)
}

// TimeFunction runs the given function fn a specified number of times
// and returns the timing results. It supports context for cancellation.
func TimeFunction(ctx context.Context, fn func(), iterations int) (TimingResult, error) {
	if iterations <= 0 {
		return TimingResult{}, errors.New(
			"iterations must be greater than 0")
	}

	times := make([]float64, 0, iterations)
	startTime := time.Now().UTC()

	// Create a progress bar with throttling
	bar := progressbar.NewOptions(iterations,
		progressbar.OptionSetDescription("Processing..."),
		progressbar.OptionShowCount(),
		progressbar.OptionShowIts(),
		progressbar.OptionSetPredictTime(true),
		progressbar.OptionClearOnFinish(),
		progressbar.OptionThrottle(100*time.Millisecond), // Update every 100ms
	)

	for i := 0; i < iterations; i++ {
		select {
		case <-ctx.Done():
			return TimingResult{}, ctx.Err()
		default:
			iterStart := time.Now().UTC()
			fn()
			elapsed := time.Since(iterStart)
			times = append(times, float64(elapsed))

			_ = bar.Add(1)
		}
	}

	endTime := time.Now().UTC()
	totalDuration := endTime.Sub(startTime)

	sort.Float64s(times)

	min := time.Duration(times[0])
	max := time.Duration(times[len(times)-1])
	mean := time.Duration(stat.Mean(times, nil))
	median := time.Duration(stat.Quantile(0.5, stat.Empirical, times, nil))
	stdDev := time.Duration(stat.StdDev(times, nil))

	return TimingResult{
		Min:        min,
		Max:        max,
		Mean:       mean,
		Median:     median,
		StdDev:     stdDev,
		StartTime:  startTime,
		EndTime:    endTime,
		Duration:   totalDuration,
		Iterations: iterations,
	}, nil
}