Enable TUI-mode pprof artifacts and trace capture

2 files changed, 147 insertions, 37 deletions

diff --git a/internal/ior.go b/internal/ior.go
index f52c584..46d0a84 100644
--- a/internal/ior.go
+++ b/internal/ior.go
@@ -8,7 +8,10 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
+	"runtime"
 	"runtime/pprof"
+	"runtime/trace"
+	"sync"
 	"syscall"
 	"time"
 
@@ -222,7 +225,7 @@ func validateRunConfig(cfg flags.Flags) error {
 }
 
 func shouldRunTraceMode(cfg flags.Flags) bool {
-	return cfg.PlainMode || cfg.FlamegraphEnable || cfg.LiveFlamegraph || cfg.PprofEnable
+	return cfg.PlainMode || cfg.FlamegraphEnable || cfg.LiveFlamegraph
 }
 
 func tuiTraceStarterFromRunTrace(
@@ -355,16 +358,88 @@ func runTraceWithContext(parentCtx context.Context, started chan<- struct{}, con
 	}
 	rb.Poll(300)
 
+	ctx := parentCtx
+	cancel := func() {}
+	if shouldAutoStopByDuration(cfg) {
+		duration := time.Duration(cfg.Duration) * time.Second
+		logln("Probing for", duration)
+		ctx, cancel = context.WithTimeout(parentCtx, duration)
+	} else {
+		logln("Probing until stopped...")
+		ctx, cancel = context.WithCancel(parentCtx)
+	}
+	defer cancel()
+
+	signalCh := make(chan os.Signal, 1)
+	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(signalCh)
+
+	go func() {
+		select {
+		case <-signalCh:
+			logln("Received signal, shutting down...")
+			cancel()
+		case <-ctx.Done():
+			return
+		}
+	}()
+
 	pprofDone := make(chan struct{})
-	var cpuProfile, memProfile *os.File
+	var cpuProfile, memProfile, execTraceProfile *os.File
+	stopExecTrace := func() {}
 	if cfg.PprofEnable {
-		if cpuProfile, err = os.Create("ior.cpuprofile"); err != nil {
+		isTUIMode := started != nil
+		cpuProfilePath, memProfilePath, execTracePath, execTraceDuration := profilingFilesForMode(isTUIMode)
+
+		if cpuProfile, err = os.Create(cpuProfilePath); err != nil {
+			return err
+		}
+		if memProfile, err = os.Create(memProfilePath); err != nil {
+			_ = cpuProfile.Close()
 			return err
 		}
-		if memProfile, err = os.Create("ior.memprofile"); err != nil {
+
+		if execTracePath != "" {
+			if execTraceProfile, err = os.Create(execTracePath); err != nil {
+				_ = cpuProfile.Close()
+				_ = memProfile.Close()
+				return err
+			}
+			if err := trace.Start(execTraceProfile); err != nil {
+				_ = cpuProfile.Close()
+				_ = memProfile.Close()
+				_ = execTraceProfile.Close()
+				return err
+			}
+
+			// TUI profiling workflow:
+			//   go tool pprof -http=:8080 ior-tui-cpu.prof
+			//   go tool trace ior-tui-trace.out
+			var stopOnce sync.Once
+			stopExecTrace = func() {
+				stopOnce.Do(func() {
+					trace.Stop()
+					_ = execTraceProfile.Close()
+				})
+			}
+
+			go func() {
+				timer := time.NewTimer(execTraceDuration)
+				defer timer.Stop()
+				select {
+				case <-ctx.Done():
+				case <-timer.C:
+				}
+				stopExecTrace()
+			}()
+		}
+
+		if err := pprof.StartCPUProfile(cpuProfile); err != nil {
+			stopExecTrace()
+			_ = cpuProfile.Close()
+			_ = memProfile.Close()
 			return err
 		}
-		pprof.StartCPUProfile(cpuProfile)
 	} else {
 		close(pprofDone)
 	}
@@ -385,31 +460,6 @@ func runTraceWithContext(parentCtx context.Context, started chan<- struct{}, con
 			origPrintCb(ep)
 		}
 	}
-	ctx := parentCtx
-	cancel := func() {}
-	if shouldAutoStopByDuration(cfg) {
-		duration := time.Duration(cfg.Duration) * time.Second
-		logln("Probing for", duration)
-		ctx, cancel = context.WithTimeout(parentCtx, duration)
-	} else {
-		logln("Probing until stopped...")
-		ctx, cancel = context.WithCancel(parentCtx)
-	}
-	defer cancel()
-
-	signalCh := make(chan os.Signal, 1)
-	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
-	defer signal.Stop(signalCh)
-
-	go func() {
-		select {
-		case <-signalCh:
-			logln("Received signal, shutting down...")
-			cancel()
-		case <-ctx.Done():
-			return
-		}
-	}()
 
 	go func() {
 		<-ctx.Done()
@@ -417,9 +467,13 @@ func runTraceWithContext(parentCtx context.Context, started chan<- struct{}, con
 			fmt.Println(el.stats())
 		}
 		if cfg.PprofEnable {
-			logln("Stoppig profiling, writing ior.cpuprofile and ior.memprofile")
+			logln("Stopping profiling and writing profile files")
 			pprof.StopCPUProfile()
-			pprof.WriteHeapProfile(memProfile)
+			runtime.GC()
+			_ = pprof.WriteHeapProfile(memProfile)
+			stopExecTrace()
+			_ = cpuProfile.Close()
+			_ = memProfile.Close()
 			close(pprofDone)
 		}
 	}()
@@ -440,5 +494,12 @@ func signalTraceStarted(started chan<- struct{}) {
 }
 
 func shouldAutoStopByDuration(cfg flags.Flags) bool {
-	return cfg.PlainMode || cfg.FlamegraphEnable || cfg.LiveFlamegraph || cfg.PprofEnable
+	return cfg.PlainMode || cfg.FlamegraphEnable || cfg.LiveFlamegraph
+}
+
+func profilingFilesForMode(tuiMode bool) (cpuProfilePath, memProfilePath, execTracePath string, execTraceDuration time.Duration) {
+	if tuiMode {
+		return "ior-tui-cpu.prof", "ior-tui-mem.prof", "ior-tui-trace.out", 10 * time.Second
+	}
+	return "ior.cpuprofile", "ior.memprofile", "", 0
 }
diff --git a/internal/ior_mode_test.go b/internal/ior_mode_test.go
index bbca555..d509cf2 100644
--- a/internal/ior_mode_test.go
+++ b/internal/ior_mode_test.go
@@ -31,8 +31,8 @@ func TestShouldRunTraceMode(t *testing.T) {
 
 	withPprof := base
 	withPprof.PprofEnable = true
-	if !shouldRunTraceMode(withPprof) {
-		t.Fatalf("expected pprof mode to use trace mode")
+	if shouldRunTraceMode(withPprof) {
+		t.Fatalf("expected pprof flag alone to keep TUI mode")
 	}
 
 	withLive := base
@@ -62,8 +62,8 @@ func TestShouldAutoStopByDuration(t *testing.T) {
 
 	withPprof := base
 	withPprof.PprofEnable = true
-	if !shouldAutoStopByDuration(withPprof) {
-		t.Fatalf("expected pprof mode to auto-stop by duration")
+	if shouldAutoStopByDuration(withPprof) {
+		t.Fatalf("expected pprof flag alone not to auto-stop by duration")
 	}
 
 	withLive := base
@@ -104,6 +104,37 @@ func TestDispatchRunUsesTraceModeWhenRequested(t *testing.T) {
 	}
 }
 
+func TestDispatchRunUsesTUIWhenOnlyPprofEnabled(t *testing.T) {
+	origRunTrace := runTraceFn
+	origRunTUI := runTUIFn
+	defer func() {
+		runTraceFn = origRunTrace
+		runTUIFn = origRunTUI
+	}()
+
+	traceCalled := false
+	tuiCalled := false
+	runTraceFn = func() error {
+		traceCalled = true
+		return nil
+	}
+	runTUIFn = func(tui.TraceStarter) error {
+		tuiCalled = true
+		return nil
+	}
+
+	cfg := flags.Flags{PprofEnable: true}
+	if err := dispatchRun(cfg); err != nil {
+		t.Fatalf("dispatchRun returned error: %v", err)
+	}
+	if traceCalled {
+		t.Fatalf("did not expect runTraceFn when only -pprof is enabled")
+	}
+	if !tuiCalled {
+		t.Fatalf("expected runTUIFn to be called")
+	}
+}
+
 func TestDispatchRunUsesTUIStarterWhenNotPlain(t *testing.T) {
 	origRunTraceWithContext := runTraceWithContextFn
 	origRunTUI := runTUIFn
@@ -220,6 +251,24 @@ func TestTuiTraceStarterFromRunTracePropagatesError(t *testing.T) {
 	}
 }
 
+func TestProfilingFilesForMode(t *testing.T) {
+	cpu, mem, execTrace, duration := profilingFilesForMode(false)
+	if cpu != "ior.cpuprofile" || mem != "ior.memprofile" {
+		t.Fatalf("unexpected trace-mode profiling file names: cpu=%q mem=%q", cpu, mem)
+	}
+	if execTrace != "" || duration != 0 {
+		t.Fatalf("expected trace-mode execution tracing to be disabled, got trace=%q duration=%s", execTrace, duration)
+	}
+
+	cpu, mem, execTrace, duration = profilingFilesForMode(true)
+	if cpu != "ior-tui-cpu.prof" || mem != "ior-tui-mem.prof" || execTrace != "ior-tui-trace.out" {
+		t.Fatalf("unexpected TUI profiling file names: cpu=%q mem=%q trace=%q", cpu, mem, execTrace)
+	}
+	if duration != 10*time.Second {
+		t.Fatalf("expected 10s TUI execution trace duration, got %s", duration)
+	}
+}
+
 func TestTuiTraceStarterFromRunTraceRespectsCancel(t *testing.T) {
 	starter := tuiTraceStarterFromRunTrace(
 		func(ctx context.Context, _ chan<- struct{}, _ func(*eventLoop)) error {