// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"runtime"
	"runtime/pprof"
	"sync/atomic"
	"syscall"
	"time"
)

func main() {
	workers := flag.Int("workers", 1, "Number of CPU-burning workers to run")
	wallTime := flag.Duration("wall-time", 2*time.Millisecond, "Wall-clock duration of each test iteration")
	iterations := flag.Int("iterations", 10, "Number of CPU profiles to collect")
	lockThread := flag.Bool("lock-thread", false, "Call runtime.LockOSThread for each worker")

	flag.Parse()

	time.Sleep(10 * time.Millisecond)
	var before, after syscall.Rusage
	err := syscall.Getrusage(syscall.RUSAGE_SELF, &before)
	if err != nil {
		log.Fatalf("syscall.Getrusagee: %v", err)
	}
	time.Sleep(10 * time.Millisecond)

	var profiles [][]byte

	for i := 0; i < *iterations; i++ {
		var cpu bytes.Buffer
		err := pprof.StartCPUProfile(&cpu)
		if err != nil {
			log.Fatalf("pprof.StartCPUProfile: %v", err)
		}

		end := time.After(*wallTime)

		var stop int32
		running := make([]int32, *workers)
		for i := range running {
			go func(running *int32) {
				if *lockThread {
					runtime.LockOSThread()
					defer runtime.UnlockOSThread()
				}
				work(running, &stop)
			}(&running[i])
		}
		for i := range running {
			for atomic.LoadInt32(&running[i]) == 0 {
				runtime.Gosched()
			}
		}

		<-end
		atomic.StoreInt32(&stop, 1)

		for i := range running {
			for atomic.LoadInt32(&running[i]) != 0 {
				runtime.Gosched()
			}
		}

		time.Sleep(10 * time.Millisecond)
		err = syscall.Getrusage(syscall.RUSAGE_SELF, &after)
		if err != nil {
			log.Fatalf("syscall.Getrusage: %v", err)
		}

		pprof.StopCPUProfile()
		profiles = append(profiles, cpu.Bytes())
	}

	utime := time.Duration(after.Utime.Nano()-before.Utime.Nano()) * time.Nanosecond
	stime := time.Duration(after.Stime.Nano()-before.Stime.Nano()) * time.Nanosecond

	buf, err := json.Marshal(struct {
		CPUProfiles [][]byte
		UTimeMicros int
		STimeMicros int
	}{
		CPUProfiles: profiles,
		UTimeMicros: int(utime.Microseconds()),
		STimeMicros: int(stime.Microseconds()),
	})
	if err != nil {
		log.Fatalf("json.Marshal: %v", err)
	}
	fmt.Printf("%s\n", buf)
}

func work(running *int32, stop *int32) {
	atomic.StoreInt32(running, 1)
	for atomic.LoadInt32(stop) == 0 {
	}
	atomic.StoreInt32(running, 0)
}
