check ameba latency

ameba_latency.go

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
	"os/signal"
	"time"

	bpf "github.com/iovisor/gobpf/bcc"
)

const source string = `
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>

struct val_t {
    u32 pid;
    char comm[TASK_COMM_LEN];
    u64 ts;
};

struct data_t {
    u32 pid;
	char comm[TASK_COMM_LEN];
	u64 delta;    
};

BPF_HASH(start, u32, struct val_t);
BPF_PERF_OUTPUT(ameba_events);

int do_entry(struct pt_regs *ctx) {
    if (!PT_REGS_PARM1(ctx))
        return 0;
    struct val_t val = {};
    u32 pid = bpf_get_current_pid_tgid();
    if (bpf_get_current_comm(&val.comm, sizeof(val.comm)) == 0) {
        val.pid = bpf_get_current_pid_tgid();
        val.ts = bpf_ktime_get_ns();
        start.update(&pid, &val);
    }

    return 0;
}

int do_return(struct pt_regs *ctx) {
    struct val_t *valp;
    struct data_t data = {};
    u64 delta;
    u32 pid = bpf_get_current_pid_tgid();
    u64 tsp = bpf_ktime_get_ns();
    valp = start.lookup(&pid);
    
	if (valp == 0)
        return 0;       // missed start
    
	bpf_probe_read(&data.comm, sizeof(data.comm), valp->comm);
    data.pid = valp->pid;
	data.delta = tsp - valp->ts;
    ameba_events.perf_submit(ctx, &data, sizeof(data));
    start.delete(&pid);
    
	return 0;
}
`

type amebaEvent struct {
	Pid uint32
	Comm [16]byte
	Delta uint64
}

func main() {
	m := bpf.NewModule(source, []string{})
	defer m.Close()

	amebaUprobe, err := m.LoadUprobe("do_entry")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load do_entry: %s\n", err)
		os.Exit(1)
	}
	amebaUretprobe, err := m.LoadUprobe("do_return")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load do_return: %s\n", err)
		os.Exit(1)
	}

	err = m.AttachUprobe("/tmp/sandbox/ameba", "ameba", amebaUprobe, -1)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to attach do_entry uprobe: %s\n", err)
		os.Exit(1)
	}
	err = m.AttachUretprobe("/tmp/sandbox/ameba", "ameba", amebaUretprobe, -1)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to attach do_return uretprobe: %s\n", err)
		os.Exit(1)
	}

	table := bpf.NewTable(m.TableId("ameba_events"), m)

	channel := make(chan []byte)

	perfMap, err := bpf.InitPerfMap(table, channel)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to init perf map: %s\n", err)
		os.Exit(1)
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	fmt.Printf("%10s\t%s\t%s\n", "PID", "COMMAND", "DURATION")
	go func() {
		var event amebaEvent
		for {
			data := <-channel
			err := binary.Read(bytes.NewBuffer(data), binary.LittleEndian, &event)
			if err != nil {
				fmt.Printf("failed to decode received data: %s\n", err)
				continue
			}
			// Convert C string (null-terminated) to Go string
			comm := string(event.Comm[:bytes.IndexByte(event.Comm[:], 0)])
			fmt.Printf("%10d\t%s\t%s\t%d\n", event.Pid, comm, time.Duration(event.Delta), event.Delta)
		}
	}()

	perfMap.Start()
	<-sig
	perfMap.Stop()
}