Unleashing the power of frame pointers for profiling pt.2 - Writing a simple profiler

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__type(key, histogram_key_t);
	__type(value, u64);
	__uint(max_entries, K_NUM_MAP_ENTRIES);
} histogram SEC(".maps");

typedef struct histogram_key {
	u32 pid;
	u32 kernel_stack_id;
	u32 user_stack_id;
} histogram_key_t;

struct {
    __uint(type, BPF_MAP_TYPE_STACK_TRACE);
    __uint(key_size, sizeof(u32));
    __uint(value_size, PERF_MAX_STACK_DEPTH * sizeof(u64));
    __uint(max_entries, 10000);
} stack_traces SEC(".maps");

SEC("perf_event")
int sample_stack_trace(struct bpf_perf_event_data* ctx)
{
	histogram_key_t key;
	// ...

	/* Sample the user and kernel stack traces, and record in the stack_traces structure. */
	key.pid = bpf_get_current_pid_tgid() >> 32;
	key.kernel_stack_id = bpf_get_stackid(ctx, &stack_traces, 0);
	key.user_stack_id = bpf_get_stackid(ctx, &stack_traces, 0 | BPF_F_USER_STACK);
	// ...
}

SEC("perf_event")
int sample_stack_trace(struct bpf_perf_event_data* ctx)
{
	histogram_key_t key;
	u64 one = 1;
	// ...

	/* Sample the user and kernel stack traces, and record in the stack_traces structure. */
	// ...

	/* Upsert stack trace histogram */
	count = (u64*)bpf_map_lookup_elem(&histogram, &key);
	if (count) {
		(*count)++;
	} else {
		bpf_map_update_elem(&histogram, &key, &one, BPF_NOEXIST);
	}

	return 0;
}

struct {
	__uint(type, BPF_MAP_TYPE_STACK_TRACE);
	__uint(key_size, sizeof(u32));
	__uint(value_size, PERF_MAX_STACK_DEPTH * sizeof(u64));
	__uint(max_entries, K_NUM_MAP_ENTRIES);
} stack_traces SEC(".maps");

import (
	"encoding/binary"
	"bytes"
	"unsafe"
)

type HistogramKey struct {
	Pid int32

	// UserStackId, an index into the stack-traces map.
	UserStackId uint32

	// KernelStackId, an index into the stack-traces map.
	KernelStackId uint32
}

// StackTrace is an array of instruction pointers (IP).
// 127 is the size of the profile, as for the default PERF_MAX_STACK_DEPTH.
type StackTrace [127]uint64

func Run(ctx context.Context) error {
	// ...
	for it := histogram.Iterator(); it.Next(); {
		k := it.Key()

		// ...
		
		var key HistogramKey
		if err = binary.Read(bytes.NewBuffer(k), binary.LittleEndian, &key); err != nil {
			// ...
		}

		var symbols string
		if int32(key.UserStackId) >= 0 {
			trace, err := getStackTrace(stackTraces, key.UserStackId)
			// ...
		}
		if int32(key.KernelStackId) >= 0 {
			trace, err := getStackTrace(stackTraces, key.KernelStackId)
			// ...
		}
		// ...
}

func getStackTrace(stackTraces *bpf.BPFMap, id uint32) (*StackTrace, error) {
	stackB, err := stackTraces.GetValue(unsafe.Pointer(&id))
	// ...

	var stackTrace StackTrace
	err = binary.Read(bytes.NewBuffer(stackB), binary.LittleEndian, &stackTrace)
	// ...

	return &stackTrace, nil
}

func calculateStats() (map[string]float64, error) {
	// ...

	traceSampleCounts := make(map[string]int, 0)
	totalSampleCount := 0
	
	// Iterate over the stack profile counts histogram map.
	for it := histogram.Iterator(); it.Next(); {
		k := it.Key()

		// ...

		// Get count for the specific sampled stack trace.
		countB, err := histogram.GetValue(unsafe.Pointer(&k[0]))
		// ...
		count := int(binary.LittleEndian.Uint64(countB))

		// ...
		
		// Increment the traceSampleCounts map value for the stack trace symbol string (e.g. "main;subfunc;")
		totalSampleCount += count
		traceSampleCounts[trace] += count
	}
	
	stats := make(map[string]float64, len(traceSampleCounts))
	for trace, count := range traceSampleCounts {
		residencyFraction := float64(count) / float64(totalSampleCount)
		stats[trace] = residencyFraction
	}

	return stats, nil
}

typedef struct {
	Elf64_Word	st_name;
	unsigned char	st_info;
	unsigned char	st_other;
	Elf64_Half	st_shndx;
	Elf64_Addr	st_value;
	Elf64_Xword	st_size;
} Elf64_Sym;

import "debug/elf"

// ELFSymTab is one of the possible abstractions around executable
// file symbol tables, for ELF files.
type ELFSymTab struct {
	symtab []elf.Symbol
}

// Load loads from the underlying filesystem the ELF file
// with debug/elf.Open and stores it in the ELFSymTab struct.
func (e *ELFSymTab) Load(pathname string) error {
	// ...
	file, err := elf.Open(pathname)
	// ...
	syms, err := file.Symbols()
	// ...
	e.symtab = syms

	return nil
}

// GetSymbol returns symbol name from an instruction pointer address
// reading the ELF symbol table.
func (e *ELFSymTab) GetSymbol(ip uint64) (string, error) {
	var sym string
	// ...
	for _, s := range e.symtab {
		if ip >= s.Value && ip < (s.Value+s.Size) {
			sym = s.Name
		}
	}

	return sym, nil
}

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__type(key, u32);			/* pid */
	__type(value, char[MAX_ARRAY_SIZE]);	/* exe_path */
	__uint(max_entries, K_NUM_MAP_ENTRIES);
} binprm_info SEC(".maps");

SEC("perf_event")
int sample_stack_trace(struct bpf_perf_event_data* ctx)
{
	// ...
	struct task_struct *task;

	/* Get current task executable pathname */
	task = (struct task_struct *)bpf_get_current_task(); /* Current task struct */
	exe_path = get_task_exe_pathname(task);
	if (exe_path == NULL) {
		return 0;
	}
	len = bpf_core_read_str(&exe_path_str, sizeof(exe_path_str), exe_path);
	if (len < 0) {
		return 0;
	}

	// ...
	/* Upsert stack trace histogram */
	count = (u64*)bpf_map_lookup_elem(&histogram, &key);
	if (count) {
		(*count)++;
	} else {
		bpf_map_update_elem(&histogram, &key, &one, BPF_NOEXIST);
		bpf_map_update_elem(&binprm_info, &key.pid, &exe_path_str, BPF_ANY);
		// ...
	}
}

SEC("perf_event")
int sample_stack_trace(struct bpf_perf_event_data* ctx)
{
	// ...
	/* Get current task executable pathname */
	task = (struct task_struct *)bpf_get_current_task(); /* Current task struct */
	exe_path = get_task_exe_pathname(task);
	// ...
}

/*
 * get_task_exe_pathname returns the task exe_file pathname.
 * This does not apply to kernel threads as they share the same memory-mapped address space,
 * as opposed to user address space.
 */
static __always_inline void *get_task_exe_pathname(struct task_struct *task)
{
	/*
	 * Get ref file path from the task's user space memory mapping descriptor.
	 * exe_file->f_path could also be accessed from current task's binprm struct 
	 * (ctx->args[2]->file->f_path)
	 */
	struct path path = BPF_CORE_READ(task, mm, exe_file, f_path);

	buffer_t *string_buf = get_buffer(0);
	if (string_buf == NULL) {
		return NULL;
	}
	/* Write path string from path struct to the buffer */
	size_t buf_off = get_pathname_from_path(&path, string_buf);
	return &string_buf->data[buf_off];
}

import (
	bpf "github.com/aquasecurity/libbpfgo"
	"github.com/pkg/errors"
	"golang.org/x/sys/unix"implements access to ELF object files.
)

func loadAndAttach(probe []byte) error {
	bpfModule, err := bpf.NewModuleFromBuffer(probe, "sample_stack_trace")
	// ...
	defer bpfModule.Close()

	if err := bpfModule.BPFLoadObject(); err != nil {
		// ...
	}

	prog, err := bpfModule.GetProgram("sample_stack_trace")
	if err != nil {
		// ...
	}

	// ...
}

import (
	bpf "github.com/aquasecurity/libbpfgo"
	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

func loadAndAttach(probe []byte) error {
	// Load the program...

	cpus := runtime.NumCPU()

	for i := 0; i < cpus; i++ {
		attr := &unix.PerfEventAttr{
			Type: unix.PERF_TYPE_SOFTWARE,		// If type is PERF_TYPE_SOFTWARE, we are measuring software events provided by the kernel.
			Config: unix.PERF_COUNT_SW_CPU_CLOCK,	// This reports the CPU clock, a high-resolution per-CPU timer.
			
			// A "sampling" event is one that generates an overflow notification every N events,
			// where N is given by sample_period.
			// sample_freq can be used if you wish to use frequency rather than period.
			// sample_period and sample_freq are mutually exclusive.
			// The kernel will adjust the sampling period to try and achieve the desired rate.
            // Prime numbers are recommended to avoid collisions with other periodical tasks.
			Sample: 11 * 1000 * 1000,
		}
		
		// Create the perf event file descriptor that corresponds to one event that is measured.
		// We're measuring a clock timer software event just to run the program on a periodic schedule.
		// When a specified number of clock samples occur, the kernel will trigger the program.
		evt, err := unix.PerfEventOpen(
			attr,	// The attribute set.
			-1,	// All the tasks.
			i,	// on the Nth CPU.
			-1,	// The group_fd argument allows event groups to be created.
			0,	// The flags.
		)
		// ...
		https://blog.px.dev/static/7b13192052f268bfd22577215d0c9f01/sample-stack-trace-function.png
		// Attach the BPF program to the sampling perf event.
		if _, err = prog.AttachPerfEvent(evt); err != nil {
			return errors.Wrap(err, "error attaching the BPF probe to the sampling perf event")
		}
	}

	return nil
}