package repro

import (
	"flag"
	"os"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

var (
	sink      int64
	taskCount = flag.Int("task-count", 100, "")
)

func init() {
	hz, _ := strconv.Atoi(os.Getenv("PROFILE_HZ"))
	if hz > 0 {
		runtime.SetCPUProfileRate(hz)
	}
}

func TestCPUProfile(t *testing.T) {
	parallelism := runtime.GOMAXPROCS(0)
	tasks := make(chan func())

	// force creation of GOMAXPROCS threads
	var gate sync.WaitGroup
	gate.Add(1)

	var wg sync.WaitGroup
	for i := 0; i < parallelism; i++ {
		gate.Add(1)
		wg.Add(1)
		go func() {
			defer wg.Done()

			runtime.LockOSThread()
			gate.Done()
			gate.Wait()
			runtime.UnlockOSThread()

			for task := range tasks {
				task()
			}
		}()
	}

	gate.Done()

	var taskWg sync.WaitGroup
	enqueueTask := func(fn func(x int) int) {
		taskWg.Add(1)
		tasks <- func() {
			defer taskWg.Done()
			var v int
			cpuHogger(fn, &v, 0)
			atomic.StoreInt64(&sink, int64(v))
		}
	}

	// do work sequentially
	for i := 0; i < *taskCount; i++ {
		enqueueTask(cpuHog1)
		taskWg.Wait()
	}

	// do work in parallel
	for i := 0; i < *taskCount; i++ {
		enqueueTask(cpuHog2)
	}
	taskWg.Wait()

	close(tasks)
	wg.Wait()
}

// cpuHogger, cpuHog1, and cpuHog2 from src/runtime/pprof/pprof_test.go

func cpuHogger(f func(x int) int, y *int, dur time.Duration) {
	// We only need to get one 100 Hz clock tick, so we've got
	// a large safety buffer.
	// But do at least 500 iterations (which should take about 100ms),
	// otherwise TestCPUProfileMultithreaded can fail if only one
	// thread is scheduled during the testing period.
	t0 := time.Now()
	accum := *y
	for i := 0; i < 500 || time.Since(t0) < dur; i++ {
		accum = f(accum)
	}
	*y = accum
}

// The actual CPU hogging function.
// Must not call other functions nor access heap/globals in the loop,
// otherwise under race detector the samples will be in the race runtime.
func cpuHog1(x int) int {
	foo := x
	for i := 0; i < 1e5; i++ {
		if foo > 0 {
			foo *= foo
		} else {
			foo *= foo + 1
		}
	}
	return foo
}

func cpuHog2(x int) int {
	foo := x
	for i := 0; i < 1e5; i++ {
		if foo > 0 {
			foo *= foo
		} else {
			foo *= foo + 2
		}
	}
	return foo
}
