Fix pprof parsing when there's more than 2 columns (#53)

Previously, the pprof parser assumed there were 2 columns, and stored
them as "samples" and "total". Instead parse all values and store
columns which we can use in future.

pprof/parser.go

@@ -51,6 +51,7 @@ type rawParser struct {
 
 	state     readState
 	funcNames map[funcID]string
+	columns   []string
 	records   []*stackRecord
 }
 
@@ -106,6 +107,7 @@ func (p *rawParser) processLine(line string) {
 			return
 		}
 	case samplesHeader:
+		p.columns = strings.Split(line, " ")
 		p.state = samples
 	case samples:
 		if strings.HasPrefix(line, "Locations") {
@@ -132,13 +134,13 @@ func (p *rawParser) toSamples() []*stack.Sample {
 		funcKey := strings.Join(funcNames, ";")
 
 		if sample, ok := samples[funcKey]; ok {
-			sample.Count += r.samples
+			sample.Count += r.samples[0]
 			continue
 		}
 
 		samples[funcKey] = &stack.Sample{
 			Funcs: funcNames,
-			Count: r.samples,
+			Count: r.samples[0],
 		}
 	}
 
@@ -171,37 +173,42 @@ func (p *rawParser) addLocation(line string) {
 }
 
 type stackRecord struct {
-	samples int64
-	total   int64
+	samples []int64
 	stack   []funcID
 }
 
 // addSample parses a sample that looks like:
 //   1   10000000: 1 2 3 4
 // and creates a stackRecord for it.
 func (p *rawParser) addSample(line string) {
-	// Parse a sample which looks like:
-	parts := splitBySpace(line)
-	if strings.HasPrefix(parts[0], "bytes:[") {
-		// Memory profiles have a line line bytes:[size1] which says the size of the object.
-		// Skip these lines.
+	if strings.Contains(line, "bytes:[") {
+		// Memory profiles have a line of bytes:[size1] which contains the size
+		// of the object. Skip these lines as we do not use it currently.
 		return
 	}
-	if len(parts) < 3 {
+
+	// Split by ":" which separates the data from the function IDs.
+	lineParts := strings.Split(line, ":")
+	if len(lineParts) != 2 {
 		p.setError(fmt.Errorf("malformed sample line: %v", line))
 		return
 	}
 
-	record := &stackRecord{
-		samples: p.parseInt(parts[0]),
-		total:   p.parseInt(strings.TrimSuffix(parts[1], ":")),
-	}
-	for _, fIDStr := range parts[2:] {
-		record.stack = append(record.stack, p.toFuncID(fIDStr))
+	samples := p.parseInts(lineParts[0])
+	funcIDs := p.parseFuncIDs(lineParts[1])
+
+	if len(samples) != len(p.columns) {
+		p.setError(fmt.Errorf("line has a different sample count (%v) than columns (%v): %v",
+			len(samples), len(p.columns), line))
+		return
 	}
 
-	p.records = append(p.records, record)
+	p.records = append(p.records, &stackRecord{
+		samples: samples,
+		stack:   funcIDs,
+	})
 }
+
 func getFunctionName(funcNames map[funcID]string, funcID funcID) string {
 	if funcName, ok := funcNames[funcID]; ok {
 		return funcName
@@ -220,6 +227,24 @@ func (r *stackRecord) funcNames(funcNames map[funcID]string) []string {
 	return names
 }
 
+func (p *rawParser) parseFuncIDs(s string) []funcID {
+	funcInts := p.parseInts(s)
+	funcIDs := make([]funcID, len(funcInts))
+	for i, fID := range funcInts {
+		funcIDs[i] = funcID(fID)
+	}
+	return funcIDs
+}
+
+func (p *rawParser) parseInts(s string) []int64 {
+	ss := splitBySpace(s)
+	samples := make([]int64, len(ss))
+	for i, s := range ss {
+		samples[i] = p.parseInt(s)
+	}
+	return samples
+}
+
 // parseInt converts a string to an int64. It stores any errors using setError.
 func (p *rawParser) parseInt(s string) int64 {
 	v, err := strconv.ParseInt(s, 10, 64)

pprof/parser_test.go

@@ -27,6 +27,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/uber/go-torch/stack"
 )
 
@@ -50,21 +51,33 @@ func parseTest1(t *testing.T) ([]byte, *rawParser) {
 
 func TestParseMemProfile(t *testing.T) {
 	parseTestRawData(t, "testdata/pprof3.raw.txt")
+	parseTestRawData(t, "testdata/pprof-memprofile-1.8.raw.txt")
 }
 
 func TestParse(t *testing.T) {
 	_, parser := parseTest1(t)
 
+	assert.Equal(t, []string{"samples/count", "cpu/nanoseconds"}, parser.columns)
+
 	// line 7 - 249 are stack records in the test file.
 	const expectedNumRecords = 242
 	if len(parser.records) != expectedNumRecords {
 		t.Errorf("Failed to parse all records, got %v records, expected %v",
 			len(parser.records), expectedNumRecords)
 	}
 	expectedRecords := map[int]*stackRecord{
-		0:  &stackRecord{1, 10000000, []funcID{1, 2, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 2, 3, 3, 3, 2, 3, 2, 3, 2, 2, 3, 2, 2, 3, 4, 5, 6}},
-		18: &stackRecord{1, 10000000, []funcID{14, 2, 2, 3, 2, 2, 3, 2, 2, 3, 3, 3, 2, 2, 2, 3, 3, 2, 3, 3, 3, 3, 3, 2, 4, 5, 6}},
-		45: &stackRecord{12, 120000000, []funcID{23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34}},
+		0: &stackRecord{
+			samples: []int64{1, 10000000},
+			stack:   []funcID{1, 2, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 2, 3, 3, 3, 2, 3, 2, 3, 2, 2, 3, 2, 2, 3, 4, 5, 6},
+		},
+		18: &stackRecord{
+			samples: []int64{1, 10000000},
+			stack:   []funcID{14, 2, 2, 3, 2, 2, 3, 2, 2, 3, 3, 3, 2, 2, 2, 3, 3, 2, 3, 3, 3, 3, 3, 2, 4, 5, 6},
+		},
+		45: &stackRecord{
+			samples: []int64{12, 120000000},
+			stack:   []funcID{23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34},
+		},
 	}
 	for recordNum, expected := range expectedRecords {
 		if got := parser.records[recordNum]; !reflect.DeepEqual(got, expected) {
@@ -150,6 +163,18 @@ func TestParseMissingSourceLine(t *testing.T) {
 	}
 }
 
+func TestParseSampleCountMismatch(t *testing.T) {
+	contents := `Samples:
+	samples/count cpu/nanoseconds alloc_objects/count
+	   2   10000000: 1
+	Locations:
+	   1: 0xaaaaa funcName :0 s=0
+`
+	_, err := ParseRaw([]byte(contents))
+	require.Error(t, err, "Expected parseRaw to fail with sample count mismatch")
+	assert.Contains(t, err.Error(), "different sample count (2) than columns (3)")
+}
+
 func testParseRawBad(t *testing.T, errorReason, errorSubstr, contents string) {
 	_, err := ParseRaw([]byte(contents))
 	if err == nil {