// Copyright 2015 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.

// Perf2pprof converts Linux perf profile data into Go pprof format.
// (The Linux perf suite is sometimes referred to as perf_events.)
//
// 	usage: perf2pprof [-list] [-o outfile] [-x exe] perf.data
//
// Perf2pprof converts perf.data, generated by the Linux perf suite,
// into pprof format. The default output file name is the input file
// name (including any directory) with .pprof appended.
// The -o flag instructs perf2pprof to write to outfile instead.
//
// Perf is a system-wide profiler for Linux: it captures data from every binary
// running on the system, along with the operating system kernel.
// Although in principle it should be possible to express those multiple
// event sources in a single pprof output file, the current perf2pprof
// is limited to extracting the samples from just one executable that
// was running. The -x flag specifies the executable, typically a full path.
// The -list flag causes perf2pprof to list all the executables with
// samples in the profile. One of -list or -x must be specified.
//
// Go and Perf
//
// By default, Go does not maintain x86 frame pointers, which means
// that perf cannot sample from Go's execution reliably.
// To build a Go 1.5 or later toolchain that works with perf, use:
//
//	GOEXPERIMENT=framepointer ./make.bash
//
// Bugs
//
// This is a very rough sketch. It does seem to work for simple perf profiles
// tracking only a single event, but it could be made much richer.
//
// It should not be necessary to specify the -x flag, as explained above.
// Even if limited to one executable's samples, perf2pprof could infer the
// correct setting for -x by reading the profile to count samples per executable,
// set -x to the executable with the most samples, and read the profile again.
//
package main // import "code.justin.tv/rhys/nursery/cmd/perf2pprof"

import (
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"os"

	"code.justin.tv/rhys/nursery/internal/_vendor/rsc.io/perf2pprof/pprof"
	"github.com/aclements/go-perf/perffile"
)

var (
	listBinaries = flag.Bool("list", false, "list executables found in profile")
	outFlag      = flag.String("o", "", "output `file` name")
	targetBinary = flag.String("x", "", "include only samples from the named `executable`")

	start = flag.Int64("start", 0, "Start timestamp (in perf's time units)")
	end   = flag.Int64("end", 0, "End timestamp, exclusive (in perf's time units)")

	vmlinuxOffset = flag.Uint64("vmlinux-offset", 0, "Kernel ASLR offset")
)

func usage() {
	fmt.Fprintf(os.Stderr, "usage: perf2pprof [-list] [-o outfile] [-x exe] [-start time] [-end time] perf.data\n")
	flag.PrintDefaults()
	os.Exit(2)
}

var numRecords int

func main() {
	log.SetPrefix("perf2pprof: ")
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 1 {
		usage()
	}

	if !*listBinaries && *targetBinary == "" {
		fmt.Fprintf(os.Stderr, "-list or -x flag is required for now\n")
		os.Exit(2)
	}

	infile := flag.Arg(0)
	outfile := *outFlag
	if outfile == "" {
		outfile = infile + ".pprof"
	}
	perf, err := perffile.Open(infile)
	if err != nil {
		log.Fatal(err)
	}

	var w *os.File
	if !*listBinaries {
		w, err = os.Create(outfile)
		if err != nil {
			log.Fatal(err)
		}
	}

	printed := make(map[string]bool)
	rs := perf.Records(perffile.RecordsFileOrder)
	unknowns := make(map[[2]string]int)
	for rs.Next() {
		if numRecords++; numRecords&(numRecords-1) == 0 {
			fmt.Printf("%d records\n", numRecords)
		}
		switch r := rs.Record.(type) {
		case *perffile.RecordThrottle, *perffile.RecordComm, *perffile.RecordExit:
			// ignore
		case *perffile.RecordFork:
			mmap[r.PID] = append([]*perffile.RecordMmap(nil), mmap[r.PPID]...)
		case *perffile.RecordMmap:
			if *listBinaries {
				if !printed[r.Filename] {
					printed[r.Filename] = true
					fmt.Printf("%q\n", r.Filename)
				}
				continue
			}
			rr := *r
			mmap[r.PID] = append(mmap[r.PID], &rr)
		case *perffile.RecordSample:
			if *listBinaries {
				continue
			}

			if *start != 0 && r.Time < uint64(*start) {
				continue
			}
			if *end != 0 && r.Time >= uint64(*end) {
				continue
			}

			switch {
			default:
				fmt.Printf("bad sample: %+v\n", *r)
				continue
			case r.Format&perffile.SampleFormatCallchain != 0:
				addSample(r.PID, r.Callchain)
			case r.Format&perffile.SampleFormatIP != 0:
				addSample(r.PID, []uint64{r.IP})
			}
		default:
			key := [2]string{fmt.Sprintf("%T", rs.Record), rs.Record.Type().String()}
			unknowns[key]++
			if n := unknowns[key]; n&(n-1) == 0 {
				fmt.Printf("found %d unknown records %s, %s\n", n, key[0], key[1])
			}
		}
	}
	if err := rs.Err(); err != nil {
		log.Fatalf("reading %s: %v", infile, err)
	}

	if *listBinaries {
		return
	}

	fmt.Printf("%d samples for %d records\n", len(samples), numRecords)

	// TODO(rsc): Get actual event type from profile.
	// TODO(rsc): A profile can actually have multiple events in it.
	// Include the event type in the stack trace map key.
	p := new(pprof.Profile)
	p.SampleType = []*pprof.ValueType{{Type: "event", Unit: "count"}}
	p.PeriodType = p.SampleType[0]
	p.Period = 1

	var stack []uint64
	var pid int

	type maploc struct {
		pc      uint64
		mapping *pprof.Mapping
	}
	locations := make(map[maploc]*pprof.Location)

	mappings := make(map[int]*pprof.Mapping)
	vmlinux := &pprof.Mapping{
		ID:     0, // This will get adjusted at the end
		File:   "vmlinux",
		Offset: 0,
		Start:  0xffffffff80000000,
		Limit:  0xffffffffffffffff,
	}
	vmlinux.Start += *vmlinuxOffset

	for _, s := range samples {
		pid, stack = decode(s.key, stack)
		_ = pid
		locs := make([]*pprof.Location, 0, len(stack))

		var mapping *pprof.Mapping
		mem := mmap[pid]
		for _, m := range mem {
			if m.Filename != *targetBinary {
				continue
			}
			var ok bool
			mapping, ok = mappings[pid]
			if !ok {
				mapping = &pprof.Mapping{
					ID:     uint64(1 + len(mappings)),
					File:   m.Filename,
					Offset: m.FileOffset,
					Start:  m.Addr,
					Limit:  m.Addr + m.Len,
				}
				mappings[pid] = mapping
				p.Mapping = append(p.Mapping, mapping)
			}
			break
		}

		for _, pc := range stack {
			if len(locs) > 0 { // pprof stores caller PC with -1 added so it lands in the call instruction
				pc--
			}

			emapping := mapping
			if pc>>63 == 1 {
				emapping = vmlinux
			}
			loc := locations[maploc{pc: pc, mapping: emapping}]
			if loc == nil {
				loc = &pprof.Location{
					Address: pc,
					ID:      uint64(1 + len(locations)),
					Mapping: emapping,
				}
				locations[maploc{pc: pc, mapping: emapping}] = loc
				p.Location = append(p.Location, loc)
			}
			locs = append(locs, loc)
		}
		p.Sample = append(p.Sample, &pprof.Sample{
			Location: locs,
			Value:    []int64{s.count},
		})
	}

	vmlinux.ID = uint64(1 + len(p.Mapping))
	p.Mapping = append(p.Mapping, vmlinux)

	if err := p.Write(w); err != nil {
		log.Fatalf("writing %s: %v", outfile, err)
	}
	if err := w.Close(); err != nil {
		log.Fatalf("writing %s: %v", outfile, err)
	}
}

var (
	samples = make(map[string]sample)
	mmap    = make(map[int][]*perffile.RecordMmap)
)

type sample struct {
	key   string
	count int64
}

var buf = make([]byte, 1024)

func filterStack(pid int, stack []uint64) []uint64 {
	mem := mmap[pid]
	stk := stack[:0]
	for _, pc := range stack {
		for _, m := range mem {
			if m.Filename != *targetBinary {
				continue
			}

			// pid is has the executable of interest mapped

			// record samples from within the executable
			if m.Addr <= pc && pc-m.Addr < m.Len {
				stk = append(stk, pc)
				break
			}

			if pc == perffile.CallchainHypervisor || pc == perffile.CallchainKernel ||
				pc == perffile.CallchainUser || pc == perffile.CallchainGuest ||
				pc == perffile.CallchainGuestKernel || pc == perffile.CallchainGuestUser {
				continue
			}
			// record samples from the process of interest calling into the kernel
			if pc>>63 == 1 {
				stk = append(stk, pc)
				// break
			}
		}
	}
	return stk
}

func addSample(pid int, stack []uint64) {
	stack = filterStack(pid, stack)
	if len(stack) == 0 {
		return
	}
	for cap(buf) < (1+len(stack))*10 {
		buf = make([]byte, cap(buf)*2)
	}
	buf = buf[:10*(1+len(stack))]
	n := 0
	n += binary.PutUvarint(buf[n:], uint64(pid))
	for _, x := range stack {
		n += binary.PutUvarint(buf[n:], x)
	}

	buf = buf[:n]
	s, ok := samples[string(buf)]
	if !ok {
		s.key = string(buf)
		if len(samples)&(len(samples)-1) == 0 {
			fmt.Printf("%d samples for %d records\n", len(samples)+1, numRecords)
		}
	}
	s.count++
	samples[s.key] = s
}

func decode(s string, stack []uint64) (pid int, outstack []uint64) {
	stack = stack[:0]
	v, n := uvarint(s)
	if n <= 0 {
		log.Fatal("malformed internal encoding")
	}
	pid = int(v)
	s = s[n:]
	for len(s) > 0 {
		v, n := uvarint(s)
		if n <= 0 {
			log.Fatal("malformed internal encoding")
		}
		stack = append(stack, v)
		s = s[n:]
	}
	return pid, stack
}

func uvarint(buf string) (uint64, int) {
	var x uint64
	var s uint
	for i := 0; i < len(buf); i++ {
		b := buf[i]
		if b < 0x80 {
			if i > 9 || i == 9 && b > 1 {
				return 0, -(i + 1) // overflow
			}
			return x | uint64(b)<<s, i + 1
		}
		x |= uint64(b&0x7f) << s
		s += 7
	}
	return 0, 0
}
