From a38cda378c9428b0621a50d640d78781e435e7ff Mon Sep 17 00:00:00 2001
From: "Harper, Jason M" <jason.m.harper@intel.com>
Date: Mon, 23 Feb 2026 14:46:57 -0800
Subject: [PATCH] fix: enhance mergeSystemFolded to handle empty sample counts
 and add tests

Signed-off-by: Harper, Jason M <jason.m.harper@intel.com>
---
 cmd/flamegraph/flamegraph_tables.go      | 14 +++++-
 cmd/flamegraph/flamegraph_tables_test.go | 54 ++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 2 deletions(-)
 create mode 100644 cmd/flamegraph/flamegraph_tables_test.go

diff --git a/cmd/flamegraph/flamegraph_tables.go b/cmd/flamegraph/flamegraph_tables.go
index 1983614e..60b667da 100644
--- a/cmd/flamegraph/flamegraph_tables.go
+++ b/cmd/flamegraph/flamegraph_tables.go
@@ -283,8 +283,18 @@ func mergeSystemFolded(perfFp string, perfDwarf string) (merged string, err erro
 	}
 	fpSampleCount := fpStacks.totalSamples()
 	dwarfSampleCount := dwarfStacks.totalSamples()
-	if fpSampleCount == 0 || dwarfSampleCount == 0 {
-		err = fmt.Errorf("sample counts cannot be zero")
+	if fpSampleCount == 0 && dwarfSampleCount == 0 {
+		err = fmt.Errorf("both fp and dwarf sample counts cannot be zero")
+		return
+	}
+	if fpSampleCount == 0 {
+		slog.Warn("no fp samples; using dwarf stacks")
+		merged = dwarfStacks.dumpFolded()
+		return
+	}
+	if dwarfSampleCount == 0 {
+		slog.Warn("no dwarf samples; using fp stacks")
+		merged = fpStacks.dumpFolded()
 		return
 	}
 	fpToDwarfScalingRatio := float64(fpSampleCount) / float64(dwarfSampleCount)
diff --git a/cmd/flamegraph/flamegraph_tables_test.go b/cmd/flamegraph/flamegraph_tables_test.go
new file mode 100644
index 00000000..e08b7556
--- /dev/null
+++ b/cmd/flamegraph/flamegraph_tables_test.go
@@ -0,0 +1,54 @@
+// Copyright (C) 2021-2025 Intel Corporation
+// SPDX-License-Identifier: BSD-3-Clause
+
+package flamegraph
+
+import "testing"
+
+func parseFoldedForTest(t *testing.T, folded string) ProcessStacks {
+	t.Helper()
+	stacks := make(ProcessStacks)
+	if err := stacks.parsePerfFolded(folded); err != nil {
+		t.Fatalf("parse failed: %v", err)
+	}
+	return stacks
+}
+
+func TestMergeSystemFoldedUsesDwarfWhenFpEmpty(t *testing.T) {
+	fpFolded := ""
+	dwarfFolded := "procA;foo;bar 3\nprocB;baz 2\n"
+
+	merged, err := mergeSystemFolded(fpFolded, dwarfFolded)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	got := parseFoldedForTest(t, merged)
+	expected := parseFoldedForTest(t, dwarfFolded)
+	if got.totalSamples() != expected.totalSamples() {
+		t.Fatalf("expected %d samples, got %d", expected.totalSamples(), got.totalSamples())
+	}
+}
+
+func TestMergeSystemFoldedUsesFpWhenDwarfEmpty(t *testing.T) {
+	fpFolded := "procA;foo;bar 3\nprocB;baz 2\n"
+	dwarfFolded := ""
+
+	merged, err := mergeSystemFolded(fpFolded, dwarfFolded)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	got := parseFoldedForTest(t, merged)
+	expected := parseFoldedForTest(t, fpFolded)
+	if got.totalSamples() != expected.totalSamples() {
+		t.Fatalf("expected %d samples, got %d", expected.totalSamples(), got.totalSamples())
+	}
+}
+
+func TestMergeSystemFoldedErrorsWhenBothEmpty(t *testing.T) {
+	_, err := mergeSystemFolded("", "")
+	if err == nil {
+		t.Fatal("expected error, got nil")
+	}
+}