fix: restore legacy flamegraph trace output mode

4 files changed, 99 insertions, 10 deletions

diff --git a/internal/flags/flags.go b/internal/flags/flags.go
index df4187c..9161378 100644
--- a/internal/flags/flags.go
+++ b/internal/flags/flags.go
@@ -42,14 +42,16 @@ type Config struct {
 	TracepointsToExclude []*regexp.Regexp
 
 	// Output/runtime flags
-	PlainMode       bool
-	TestFlames      bool
-	TestLiveFlames  bool
-	LiveInterval    time.Duration
-	TUIExportEnable bool
-	CollapsedFields []string
-	CountField      string
-	GlobalFilter    globalfilter.Filter
+	PlainMode        bool
+	FlamegraphOutput bool
+	OutputName       string
+	TestFlames       bool
+	TestLiveFlames   bool
+	LiveInterval     time.Duration
+	TUIExportEnable  bool
+	CollapsedFields  []string
+	CountField       string
+	GlobalFilter     globalfilter.Filter
 }
 
 // NewFlags returns a configuration instance initialized with project defaults.
@@ -167,6 +169,8 @@ func parse() error {
 	tracepointsToExclude := flag.String("tpsExclude", "", "Comma separated list regexes for tracepoints to exclude")
 
 	flag.BoolVar(&cfg.PlainMode, "plain", false, "Enable plain CSV output mode (disable TUI)")
+	flag.BoolVar(&cfg.FlamegraphOutput, "flamegraph", false, "Write aggregated .ior.zst output for trace/integration workflows")
+	flag.StringVar(&cfg.OutputName, "name", cfg.OutputName, "Base name for .ior.zst trace output files")
 	flag.BoolVar(&cfg.TestFlames, "testflames", false, "Run TUI with static synthetic flamegraph data for keyboard-navigation testing")
 	flag.BoolVar(&cfg.TestLiveFlames, "testliveflames", false, "Run TUI with continuously-updating synthetic flamegraph data for live keyboard-navigation testing")
 	flag.DurationVar(&cfg.LiveInterval, "live-interval", cfg.LiveInterval, "Synthetic live flamegraph refresh interval for --testliveflames")
diff --git a/internal/flags/flags_test.go b/internal/flags/flags_test.go
index 2469068..4485f34 100644
--- a/internal/flags/flags_test.go
+++ b/internal/flags/flags_test.go
@@ -102,6 +102,19 @@ func TestParseTestLiveFlamesFlag(t *testing.T) {
 	}
 }
 
+func TestParseFlamegraphOutputFlags(t *testing.T) {
+	cfg, err := parseForTest(t, "--flamegraph", "--name", "scenario-run")
+	if err != nil {
+		t.Fatalf("parse returned error: %v", err)
+	}
+	if !cfg.FlamegraphOutput {
+		t.Fatalf("expected --flamegraph to enable .ior.zst output mode")
+	}
+	if got, want := cfg.OutputName, "scenario-run"; got != want {
+		t.Fatalf("output name = %q, want %q", got, want)
+	}
+}
+
 func TestParseDefaultCollapsedFieldsOrder(t *testing.T) {
 	cfg, err := parseForTest(t)
 	if err != nil {
diff --git a/internal/flamegraph/recorder.go b/internal/flamegraph/recorder.go
new file mode 100644
index 0000000..432e509
--- /dev/null
+++ b/internal/flamegraph/recorder.go
@@ -0,0 +1,34 @@
+package flamegraph
+
+import "ior/internal/event"
+
+// Recorder aggregates event pairs and writes them to the legacy .ior.zst format.
+// Integration tests still use this artifact to assert trace output end-to-end.
+type Recorder struct {
+	name string
+	data iorData
+}
+
+// NewRecorder creates a recorder for one trace run.
+func NewRecorder(name string) *Recorder {
+	return &Recorder{
+		name: name,
+		data: newIorData(),
+	}
+}
+
+// AddPair folds one traced syscall pair into the aggregated output.
+func (r *Recorder) AddPair(pair *event.Pair) {
+	if r == nil || pair == nil {
+		return
+	}
+	r.data.addEventPair(pair)
+}
+
+// Write persists the aggregated trace output to a .ior.zst file.
+func (r *Recorder) Write() error {
+	if r == nil {
+		return nil
+	}
+	return r.data.serializeToFile(r.name)
+}
diff --git a/internal/ior.go b/internal/ior.go
index 3f145a9..8ad82c0 100644
--- a/internal/ior.go
+++ b/internal/ior.go
@@ -91,9 +91,18 @@ func validateRunConfig(cfg flags.Config) error {
 	if cfg.TestFlames && cfg.PlainMode {
 		return errors.New("--testflames cannot be combined with -plain")
 	}
+	if cfg.TestFlames && cfg.FlamegraphOutput {
+		return errors.New("--testflames cannot be combined with -flamegraph")
+	}
 	if cfg.TestLiveFlames && cfg.PlainMode {
 		return errors.New("--testliveflames cannot be combined with -plain")
 	}
+	if cfg.TestLiveFlames && cfg.FlamegraphOutput {
+		return errors.New("--testliveflames cannot be combined with -flamegraph")
+	}
+	if cfg.PlainMode && cfg.FlamegraphOutput {
+		return errors.New("-plain and -flamegraph are mutually exclusive")
+	}
 	if cfg.TestFlames && cfg.TestLiveFlames {
 		return errors.New("--testflames and --testliveflames are mutually exclusive")
 	}
@@ -168,7 +177,7 @@ func runSyntheticLiveFlames(ctx context.Context, liveTrie *flamegraph.LiveTrie,
 }
 
 func shouldRunTraceMode(cfg flags.Config) bool {
-	return cfg.PlainMode
+	return cfg.PlainMode || cfg.FlamegraphOutput
 }
 
 func tuiTraceStarterFromRunTrace(
@@ -525,6 +534,10 @@ func runTraceWithContext(parentCtx context.Context, cfg flags.Config, started ch
 
 	verbose := started == nil
 	logln := newLogger(verbose)
+	var recorder *flamegraph.Recorder
+	if cfg.FlamegraphOutput {
+		recorder = flamegraph.NewRecorder(cfg.OutputName)
+	}
 
 	bpfModule, mgr, releaseBindings, err := setupBPFModule(parentCtx, cfg)
 	if err != nil {
@@ -553,6 +566,15 @@ func runTraceWithContext(parentCtx context.Context, cfg flags.Config, started ch
 	if err != nil {
 		return err
 	}
+	if recorder != nil {
+		recordOutput := func(el *eventLoop) {
+			el.printCb = func(ep *event.Pair) {
+				recorder.AddPair(ep)
+				ep.Recycle()
+			}
+		}
+		configure = chainEventLoopConfigure(recordOutput, configure)
+	}
 	configureEventLoopOutput(el, mgr, configure)
 	startTraceShutdownWatcher(ctx, verbose, el, profiling, logln)
 
@@ -560,10 +582,26 @@ func runTraceWithContext(parentCtx context.Context, cfg flags.Config, started ch
 	el.run(ctx, ch)
 	totalDuration := time.Since(startTime)
 	<-profiling.done
+	if recorder != nil {
+		if err := recorder.Write(); err != nil {
+			return err
+		}
+	}
 	logln("Good bye... (unloading BPF tracepoints will take a few seconds...) after", totalDuration)
 	return nil
 }
 
+func chainEventLoopConfigure(fns ...func(*eventLoop)) func(*eventLoop) {
+	return func(el *eventLoop) {
+		for _, fn := range fns {
+			if fn == nil {
+				continue
+			}
+			fn(el)
+		}
+	}
+}
+
 func signalTraceStarted(started chan<- struct{}) {
 	if started == nil {
 		return
@@ -572,7 +610,7 @@ func signalTraceStarted(started chan<- struct{}) {
 }
 
 func shouldAutoStopByDuration(cfg flags.Config) bool {
-	return cfg.PlainMode
+	return cfg.PlainMode || cfg.FlamegraphOutput
 }
 
 func profilingFilesForMode(tuiMode bool) (cpuProfilePath, memProfilePath, execTracePath string, execTraceDuration time.Duration) {