LuaJIT-to-Go/bench/bench_profile.go

package luajit_bench

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"
	"testing"
)

// Profiling flags
var (
	cpuProfile   = flag.String("cpuprofile", "", "write cpu profile to `file`")
	memProfile   = flag.String("memprofile", "", "write memory profile to `file`")
	memProfileGC = flag.Bool("memprofilegc", false, "force GC before writing memory profile")
	blockProfile = flag.String("blockprofile", "", "write block profile to `file`")
	mutexProfile = flag.String("mutexprofile", "", "write mutex profile to `file`")
)

// setupTestMain configures profiling for benchmarks
func setupTestMain() {
	// Make sure the flags are parsed
	if !flag.Parsed() {
		flag.Parse()
	}

	// CPU profiling
	if *cpuProfile != "" {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to create CPU profile: %v\n", err)
			os.Exit(1)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to start CPU profile: %v\n", err)
			os.Exit(1)
		}
		fmt.Println("CPU profiling enabled")
	}

	// Block profiling (goroutine blocking)
	if *blockProfile != "" {
		runtime.SetBlockProfileRate(1)
		fmt.Println("Block profiling enabled")
	}

	// Mutex profiling (lock contention)
	if *mutexProfile != "" {
		runtime.SetMutexProfileFraction(1)
		fmt.Println("Mutex profiling enabled")
	}
}

// teardownTestMain completes profiling and writes output files
func teardownTestMain() {
	// Stop CPU profile
	if *cpuProfile != "" {
		pprof.StopCPUProfile()
		fmt.Println("CPU profile written to", *cpuProfile)
	}

	// Write memory profile
	if *memProfile != "" {
		f, err := os.Create(*memProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to create memory profile: %v\n", err)
			os.Exit(1)
		}
		defer f.Close()

		// Force garbage collection before writing memory profile if requested
		if *memProfileGC {
			runtime.GC()
		}

		if err := pprof.WriteHeapProfile(f); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to write memory profile: %v\n", err)
			os.Exit(1)
		}
		fmt.Println("Memory profile written to", *memProfile)
	}

	// Write block profile
	if *blockProfile != "" {
		f, err := os.Create(*blockProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to create block profile: %v\n", err)
			os.Exit(1)
		}
		defer f.Close()

		if err := pprof.Lookup("block").WriteTo(f, 0); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to write block profile: %v\n", err)
			os.Exit(1)
		}
		fmt.Println("Block profile written to", *blockProfile)
	}

	// Write mutex profile
	if *mutexProfile != "" {
		f, err := os.Create(*mutexProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to create mutex profile: %v\n", err)
			os.Exit(1)
		}
		defer f.Close()

		if err := pprof.Lookup("mutex").WriteTo(f, 0); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to write mutex profile: %v\n", err)
			os.Exit(1)
		}
		fmt.Println("Mutex profile written to", *mutexProfile)
	}
}

// TestMain is the entry point for all tests in this package
func TestMain(m *testing.M) {
	setupTestMain()
	code := m.Run()
	teardownTestMain()
	os.Exit(code)
}

// MemStats captures a snapshot of memory statistics
type MemStats struct {
	Alloc      uint64
	TotalAlloc uint64
	Sys        uint64
	Mallocs    uint64
	Frees      uint64
	HeapAlloc  uint64
}

// CaptureMemStats returns current memory statistics
func CaptureMemStats() MemStats {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	return MemStats{
		Alloc:      m.Alloc,
		TotalAlloc: m.TotalAlloc,
		Sys:        m.Sys,
		Mallocs:    m.Mallocs,
		Frees:      m.Frees,
		HeapAlloc:  m.HeapAlloc,
	}
}

// TrackMemoryUsage runs fn and reports memory usage before and after
func TrackMemoryUsage(b *testing.B, name string, fn func()) {
	b.Helper()

	// Force GC before measurement
	runtime.GC()

	// Capture memory stats before
	before := CaptureMemStats()

	// Run the function
	fn()

	// Force GC after measurement to get accurate stats
	runtime.GC()

	// Capture memory stats after
	after := CaptureMemStats()

	// Report stats
	b.ReportMetric(float64(after.Mallocs-before.Mallocs), name+"-mallocs")
	b.ReportMetric(float64(after.TotalAlloc-before.TotalAlloc)/float64(b.N), name+"-bytes/op")
}