internal/pprof: a function to compute total pprof time

Also, a test.

Change-Id: I86c777a7519ba5cf6c9980eb2e7ff3acdba4031f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/507885
Auto-Submit: Alan Donovan <[email protected]>
Run-TryBot: Alan Donovan <[email protected]>
Reviewed-by: Robert Findley <[email protected]>
gopls-CI: kokoro <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>

internal/pprof/main.go

@@ -0,0 +1,36 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build ignore
+// +build ignore
+
+// The pprof command prints the total time in a pprof profile provided
+// through the standard input.
+package main
+
+import (
+	"compress/gzip"
+	"fmt"
+	"io"
+	"log"
+	"os"
+
+	"golang.org/x/tools/internal/pprof"
+)
+
+func main() {
+	rd, err := gzip.NewReader(os.Stdin)
+	if err != nil {
+		log.Fatal(err)
+	}
+	payload, err := io.ReadAll(rd)
+	if err != nil {
+		log.Fatal(err)
+	}
+	total, err := pprof.TotalTime(payload)
+	if err != nil {
+		log.Fatal(err)
+	}
+	fmt.Println(total)
+}

internal/pprof/pprof.go

@@ -0,0 +1,89 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package pprof provides minimalistic routines for extracting
+// information from profiles.
+package pprof
+
+import (
+	"fmt"
+	"time"
+)
+
+// TotalTime parses the profile data and returns the accumulated time.
+// The input should not be gzipped.
+func TotalTime(data []byte) (total time.Duration, err error) {
+	defer func() {
+		if x := recover(); x != nil {
+			err = fmt.Errorf("error parsing pprof profile: %v", x)
+		}
+	}()
+	decode(&total, data, msgProfile)
+	return
+}
+
+// All errors are handled by panicking.
+// Constants are copied below to avoid dependency on protobufs or pprof.
+
+// protobuf wire types, from https://developers.google.com/protocol-buffers/docs/encoding
+const (
+	wireVarint = 0
+	wireBytes  = 2
+)
+
+// pprof field numbers, from https://github.com/google/pprof/blob/master/proto/profile.proto
+const (
+	fldProfileSample = 2 // repeated Sample
+	fldSampleValue   = 2 // repeated int64
+)
+
+// arbitrary numbering of message types
+const (
+	msgProfile = 0
+	msgSample  = 1
+)
+
+func decode(total *time.Duration, data []byte, msg int) {
+	for len(data) > 0 {
+		// Read tag (wire type and field number).
+		tag := varint(&data)
+
+		// Read wire value (int or bytes).
+		wire := tag & 7
+		var ival uint64
+		var sval []byte
+		switch wire {
+		case wireVarint:
+			ival = varint(&data)
+
+		case wireBytes:
+			n := varint(&data)
+			sval, data = data[:n], data[n:]
+
+		default:
+			panic(fmt.Sprintf("unexpected wire type: %d", wire))
+		}
+
+		// Process field of msg.
+		fld := tag >> 3
+		switch {
+		case msg == msgProfile && fld == fldProfileSample:
+			decode(total, sval, msgSample) // recursively decode Sample message
+
+		case msg == msgSample, fld == fldSampleValue:
+			*total += time.Duration(ival) // accumulate time
+		}
+	}
+}
+
+func varint(data *[]byte) (v uint64) {
+	for i := 0; ; i++ {
+		b := uint64((*data)[i])
+		v += (b & 0x7f) << (7 * i)
+		if b < 0x80 {
+			*data = (*data)[i+1:]
+			return v
+		}
+	}
+}

internal/pprof/pprof_test.go

@@ -0,0 +1,46 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package pprof_test
+
+import (
+	"bytes"
+	"compress/gzip"
+	"io"
+	"log"
+	"os"
+	"testing"
+	"time"
+
+	"golang.org/x/tools/internal/pprof"
+)
+
+func TestTotalTime(t *testing.T) {
+	// $ go tool pprof testdata/sample.pprof <&- 2>&1 | grep Total
+	// Duration: 11.10s, Total samples = 27.59s (248.65%)
+	const (
+		filename = "testdata/sample.pprof"
+		want     = time.Duration(27590003550)
+	)
+
+	profGz, err := os.ReadFile(filename)
+	if err != nil {
+		t.Fatal(err)
+	}
+	rd, err := gzip.NewReader(bytes.NewReader(profGz))
+	if err != nil {
+		t.Fatal(err)
+	}
+	payload, err := io.ReadAll(rd)
+	if err != nil {
+		t.Fatal(err)
+	}
+	got, err := pprof.TotalTime(payload)
+	if err != nil {
+		log.Fatal(err)
+	}
+	if got != want {
+		t.Fatalf("TotalTime(%q): got %v (%d), want %v (%d)", filename, got, got, want, want)
+	}
+}