feat: Propate the uprobe link into sample type

With #651 it is possible to attach uprobes, but they currently are
merged together. This PR introduces the use of cookies (available from kernel 5.15+) to keep
track of the links that lead to the events being generated.

It now uses the symbol part of the uprobe link and propages it into the sample type. That is something that we currently don't do well enough.

So let's look at this example:

```
--uprobe-link /nix/store/r7pnxs3cfl3qxwacj38iakpm5v1ch6lz-glibc-2.40-66/lib/libc.so.6:malloc --uprobe-link /nix/store/r7pnxs3cfl3qxwacj38iakpm5v1ch6lz-glibc-2.40-66/lib/libc.so.6:open --uprobe-link /proc/4501/root/usr/lib/aarch64-linux-gnu/libc.so.6:open --uprobe-link /proc/4501/root/usr/lib/aarch64-linux-gnu/libc.so.6:malloc
```

There would be resulting two sample types

- `uprobe_malloc_events:count`
- `uprobe_open_events:count`

Author: simonswine

cli_flags.go

@@ -13,6 +13,7 @@ import (
 
 	"go.opentelemetry.io/ebpf-profiler/internal/controller"
 	"go.opentelemetry.io/ebpf-profiler/tracer"
+	"go.opentelemetry.io/ebpf-profiler/tracer/types"
 )
 
 const (
@@ -131,7 +132,11 @@ func parseArgs() (*controller.Config, error) {
 
 	fs.StringVar(&args.IncludeEnvVars, "env-vars", defaultEnvVarsValue, envVarsHelp)
 
-	fs.Func("uprobe-link", probeLinkHelper, func(link string) error {
+	fs.Func("uprobe-link", probeLinkHelper, func(linkStr string) error {
+		link, err := types.ParseUProbeLink(linkStr)
+		if err != nil {
+			return err
+		}
 		args.UProbeLinks = append(args.UProbeLinks, link)
 		return nil
 	})

internal/controller/config.go

@@ -11,6 +11,7 @@ import (
 
 	"go.opentelemetry.io/ebpf-profiler/reporter"
 	"go.opentelemetry.io/ebpf-profiler/tracer"
+	"go.opentelemetry.io/ebpf-profiler/tracer/types"
 )
 
 type Config struct {
@@ -32,7 +33,7 @@ type Config struct {
 	VerboseMode            bool
 	Version                bool
 	OffCPUThreshold        float64
-	UProbeLinks            []string
+	UProbeLinks            []types.UProbeLink
 	LoadProbe              bool
 
 	Reporter reporter.Reporter

main.go

@@ -117,6 +117,7 @@ func mainWithExitCode() exitCode {
 		GRPCConnectionTimeout:  intervals.GRPCConnectionTimeout(),
 		ReportInterval:         intervals.ReportInterval(),
 		SamplesPerSecond:       cfg.SamplesPerSecond,
+		UProbeLinks:            cfg.UProbeLinks,
 	})
 	if err != nil {
 		log.Error(err)

reporter/base_reporter.go

@@ -41,10 +41,18 @@ func (b *baseReporter) Stop() {
 }
 
 func (b *baseReporter) ReportTraceEvent(trace *libpf.Trace, meta *samples.TraceEventMeta) error {
-	switch meta.Origin {
+	origin := samples.Origin{
+		Origin: meta.Origin,
+	}
+	switch origin.Origin {
 	case support.TraceOriginSampling:
 	case support.TraceOriginOffCPU:
 	case support.TraceOriginUProbe:
+		uprobeIdx := int(meta.OffTime - 1)
+		meta.OffTime = 0
+		if uprobeIdx >= 0 && uprobeIdx < len(b.cfg.UProbeLinks) {
+			origin.ProbeLinkName = b.cfg.UProbeLinks[uprobeIdx].Symbol
+		}
 	default:
 		return fmt.Errorf("skip reporting trace for %d origin: %w", meta.Origin,
 			errUnknownOrigin)
@@ -72,21 +80,21 @@ func (b *baseReporter) ReportTraceEvent(trace *libpf.Trace, meta *samples.TraceE
 
 	if _, exists := (*eventsTree)[samples.ContainerID(containerID)]; !exists {
 		(*eventsTree)[samples.ContainerID(containerID)] =
-			make(map[libpf.Origin]samples.KeyToEventMapping)
+			make(map[samples.Origin]samples.KeyToEventMapping)
 	}
 
-	if _, exists := (*eventsTree)[samples.ContainerID(containerID)][meta.Origin]; !exists {
-		(*eventsTree)[samples.ContainerID(containerID)][meta.Origin] =
+	if _, exists := (*eventsTree)[samples.ContainerID(containerID)][origin]; !exists {
+		(*eventsTree)[samples.ContainerID(containerID)][origin] =
 			make(samples.KeyToEventMapping)
 	}
 
-	if events, exists := (*eventsTree)[samples.ContainerID(containerID)][meta.Origin][key]; exists {
+	if events, exists := (*eventsTree)[samples.ContainerID(containerID)][origin][key]; exists {
 		events.Timestamps = append(events.Timestamps, uint64(meta.Timestamp))
 		events.OffTimes = append(events.OffTimes, meta.OffTime)
-		(*eventsTree)[samples.ContainerID(containerID)][meta.Origin][key] = events
+		(*eventsTree)[samples.ContainerID(containerID)][origin][key] = events
 		return nil
 	}
-	(*eventsTree)[samples.ContainerID(containerID)][meta.Origin][key] = &samples.TraceEvents{
+	(*eventsTree)[samples.ContainerID(containerID)][origin][key] = &samples.TraceEvents{
 		Frames:     trace.Frames,
 		Timestamps: []uint64{uint64(meta.Timestamp)},
 		OffTimes:   []int64{meta.OffTime},

reporter/config.go

@@ -9,6 +9,7 @@ import (
 	"google.golang.org/grpc"
 
 	"go.opentelemetry.io/ebpf-profiler/reporter/samples"
+	"go.opentelemetry.io/ebpf-profiler/tracer/types"
 )
 
 type Config struct {
@@ -47,4 +48,6 @@ type Config struct {
 	// GRPCDialOptions allows passing additional gRPC dial options when establishing
 	// the connection to the collector. These options are appended after the default options.
 	GRPCDialOptions []grpc.DialOption
+
+	UProbeLinks []types.UProbeLink
 }

reporter/internal/pdata/generate.go

@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"math"
 	"path/filepath"
+	"slices"
+	"strings"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -79,11 +81,19 @@ func (p *Pdata) Generate(tree samples.TraceEventsTree,
 		sp.Scope().SetVersion(agentVersion)
 		sp.SetSchemaUrl(semconv.SchemaURL)
 
-		for _, origin := range []libpf.Origin{
-			support.TraceOriginSampling,
-			support.TraceOriginOffCPU,
-			support.TraceOriginUProbe,
-		} {
+		origins := make([]samples.Origin, len(originToEvents))
+		for o := range originToEvents {
+			origins = append(origins, o)
+		}
+		slices.SortFunc(origins, func(a, b samples.Origin) int {
+			if v := a.Origin - b.Origin; v != 0 {
+				return int(v)
+			}
+
+			return strings.Compare(a.ProbeLinkName, b.ProbeLinkName)
+		})
+
+		for _, origin := range origins {
 			if len(originToEvents[origin]) == 0 {
 				// Do not append empty profiles.
 				continue
@@ -129,12 +139,12 @@ func (p *Pdata) setProfile(
 	mappingSet orderedset.OrderedSet[uniqueMapping],
 	stackSet orderedset.OrderedSet[stackInfo],
 	locationSet orderedset.OrderedSet[locationInfo],
-	origin libpf.Origin,
+	origin samples.Origin,
 	events map[samples.TraceAndMetaKey]*samples.TraceEvents,
 	profile pprofile.Profile,
 ) error {
 	st := profile.SampleType()
-	switch origin {
+	switch origin.Origin {
 	case support.TraceOriginSampling:
 		profile.SetPeriod(1e9 / int64(p.samplesPerSecond))
 		pt := profile.PeriodType()
@@ -147,11 +157,15 @@ func (p *Pdata) setProfile(
 		st.SetTypeStrindex(stringSet.Add("events"))
 		st.SetUnitStrindex(stringSet.Add("nanoseconds"))
 	case support.TraceOriginUProbe:
-		st.SetTypeStrindex(stringSet.Add("events"))
+		if origin.ProbeLinkName != "" {
+			st.SetTypeStrindex(stringSet.Add(fmt.Sprintf("uprobe_%s_events", origin.ProbeLinkName)))
+		} else {
+			st.SetTypeStrindex(stringSet.Add("uprobe_events"))
+		}
 		st.SetUnitStrindex(stringSet.Add("count"))
 	default:
 		// Should never happen
-		return fmt.Errorf("generating profile for unsupported origin %d", origin)
+		return fmt.Errorf("generating profile for unsupported origin %d", origin.Origin)
 	}
 
 	startTS, endTS := uint64(math.MaxUint64), uint64(0)
@@ -164,7 +178,7 @@ func (p *Pdata) setProfile(
 		}
 
 		sample.TimestampsUnixNano().FromRaw(traceInfo.Timestamps)
-		if origin == support.TraceOriginOffCPU {
+		if origin.Origin == support.TraceOriginOffCPU {
 			sample.Values().Append(traceInfo.OffTimes...)
 		}
 

reporter/internal/pdata/generate_test.go

@@ -146,21 +146,21 @@ func newTestFrames(extraFrame bool) libpf.Frames {
 func TestFunctionTableOrder(t *testing.T) {
 	for _, tt := range []struct {
 		name   string
-		events map[libpf.Origin]samples.KeyToEventMapping
+		events map[samples.Origin]samples.KeyToEventMapping
 
 		wantFunctionTable        []string
 		expectedResourceProfiles int
 	}{
 		{
 			name:                     "no events",
-			events:                   map[libpf.Origin]samples.KeyToEventMapping{},
+			events:                   map[samples.Origin]samples.KeyToEventMapping{},
 			wantFunctionTable:        []string{""},
 			expectedResourceProfiles: 0,
 		}, {
 			name:                     "single executable",
 			expectedResourceProfiles: 1,
-			events: map[libpf.Origin]samples.KeyToEventMapping{
-				support.TraceOriginSampling: map[samples.TraceAndMetaKey]*samples.TraceEvents{
+			events: map[samples.Origin]samples.KeyToEventMapping{
+				samples.Origin{Origin: support.TraceOriginSampling}: map[samples.TraceAndMetaKey]*samples.TraceEvents{
 					{Pid: 1}: {
 						Frames:     newTestFrames(false),
 						Timestamps: []uint64{1, 2, 3, 4, 5},
@@ -210,12 +210,12 @@ func TestFunctionTableOrder(t *testing.T) {
 func TestProfileDuration(t *testing.T) {
 	for _, tt := range []struct {
 		name   string
-		events map[libpf.Origin]samples.KeyToEventMapping
+		events map[samples.Origin]samples.KeyToEventMapping
 	}{
 		{
 			name: "profile duration",
-			events: map[libpf.Origin]samples.KeyToEventMapping{
-				support.TraceOriginSampling: map[samples.TraceAndMetaKey]*samples.TraceEvents{
+			events: map[samples.Origin]samples.KeyToEventMapping{
+				samples.Origin{Origin: support.TraceOriginSampling}: map[samples.TraceAndMetaKey]*samples.TraceEvents{
 					{Pid: 1}: {
 						Timestamps: []uint64{2, 1, 3, 4, 7},
 					},
@@ -286,8 +286,8 @@ func TestGenerate_SingleContainerSingleOrigin(t *testing.T) {
 		Tid:            456,
 		ApmServiceName: "svc",
 	}
-	events := map[libpf.Origin]samples.KeyToEventMapping{
-		support.TraceOriginSampling: {
+	events := map[samples.Origin]samples.KeyToEventMapping{
+		samples.Origin{Origin: support.TraceOriginSampling}: {
 			traceKey: &samples.TraceEvents{
 				Frames: singleFrameTrace(libpf.GoFrame, mappingFile,
 					0x10, funcName, filePath, 42),
@@ -351,23 +351,23 @@ func TestGenerate_MultipleOriginsAndContainers(t *testing.T) {
 	traceKey := samples.TraceAndMetaKey{ExecutablePath: "/bin/foo"}
 	frames := singleFrameTrace(libpf.PythonFrame, mappingFile, 0x20, "f", "/bin/foo", 1)
 
-	events1 := map[libpf.Origin]samples.KeyToEventMapping{
-		support.TraceOriginSampling: {
+	events1 := map[samples.Origin]samples.KeyToEventMapping{
+		samples.Origin{Origin: support.TraceOriginSampling}: {
 			traceKey: &samples.TraceEvents{
 				Frames:     frames,
 				Timestamps: []uint64{1, 2},
 			},
 		},
-		support.TraceOriginOffCPU: {
+		samples.Origin{Origin: support.TraceOriginOffCPU}: {
 			traceKey: &samples.TraceEvents{
 				Frames:     frames,
 				Timestamps: []uint64{3, 4},
 				OffTimes:   []int64{10, 20},
 			},
 		},
 	}
-	events2 := map[libpf.Origin]samples.KeyToEventMapping{
-		support.TraceOriginSampling: {
+	events2 := map[samples.Origin]samples.KeyToEventMapping{
+		samples.Origin{Origin: support.TraceOriginSampling}: {
 			traceKey: &samples.TraceEvents{
 				Frames:     frames,
 				Timestamps: []uint64{5},
@@ -412,8 +412,8 @@ func TestGenerate_StringAndFunctionTablePopulation(t *testing.T) {
 	})
 
 	traceKey := samples.TraceAndMetaKey{ExecutablePath: filePath}
-	events := map[libpf.Origin]samples.KeyToEventMapping{
-		support.TraceOriginSampling: {
+	events := map[samples.Origin]samples.KeyToEventMapping{
+		samples.Origin{Origin: support.TraceOriginSampling}: {
 			traceKey: &samples.TraceEvents{
 				Frames: singleFrameTrace(libpf.PythonFrame, mappingFile, 0x30,
 					funcName, filePath, 123),
@@ -476,8 +476,8 @@ func TestGenerate_NativeFrame(t *testing.T) {
 		Pid:            789,
 		Tid:            1011,
 	}
-	events := map[libpf.Origin]samples.KeyToEventMapping{
-		support.TraceOriginSampling: {
+	events := map[samples.Origin]samples.KeyToEventMapping{
+		samples.Origin{Origin: support.TraceOriginSampling}: {
 			traceKey: &samples.TraceEvents{
 				Frames:     singleFrameNative(mappingFile, 0x1000, 0x1000, 0x2000, 0x100),
 				Timestamps: []uint64{123, 456, 789},
@@ -554,15 +554,15 @@ func TestGenerate_NativeFrame(t *testing.T) {
 func TestStackTableOrder(t *testing.T) {
 	for _, tt := range []struct {
 		name   string
-		events map[libpf.Origin]samples.KeyToEventMapping
+		events map[samples.Origin]samples.KeyToEventMapping
 
 		wantStackTable           [][]int32
 		expectedLocationTableLen int
 	}{
 		{
 			name: "single stack",
-			events: map[libpf.Origin]samples.KeyToEventMapping{
-				support.TraceOriginSampling: map[samples.TraceAndMetaKey]*samples.TraceEvents{
+			events: map[samples.Origin]samples.KeyToEventMapping{
+				samples.Origin{Origin: support.TraceOriginSampling}: map[samples.TraceAndMetaKey]*samples.TraceEvents{
 					{}: {
 						Frames:     newTestFrames(false),
 						Timestamps: []uint64{1, 2, 3, 4, 5},
@@ -576,16 +576,16 @@ func TestStackTableOrder(t *testing.T) {
 		},
 		{
 			name: "multiple stacks",
-			events: map[libpf.Origin]samples.KeyToEventMapping{
-				support.TraceOriginSampling: map[samples.TraceAndMetaKey]*samples.TraceEvents{
+			events: map[samples.Origin]samples.KeyToEventMapping{
+				samples.Origin{Origin: support.TraceOriginSampling}: map[samples.TraceAndMetaKey]*samples.TraceEvents{
 					{Pid: 1}: {
 						Frames:     newTestFrames(false),
 						Timestamps: []uint64{1, 2, 3, 4, 5},
 					},
 				},
 				// This test relies on an implementation detail for ordering of results:
 				// it assumes that support.TraceOriginSampling events are processed first
-				support.TraceOriginOffCPU: map[samples.TraceAndMetaKey]*samples.TraceEvents{
+				samples.Origin{Origin: support.TraceOriginOffCPU}: map[samples.TraceAndMetaKey]*samples.TraceEvents{
 					{Pid: 2}: {
 						Frames:     newTestFrames(true),
 						Timestamps: []uint64{7, 8, 9, 10, 11, 12},

reporter/samples/samples.go

@@ -47,14 +47,22 @@ type TraceAndMetaKey struct {
 	// Executable path is retrieved from /proc/PID/exe
 	ExecutablePath string
 
+	// Name of the hook that collected this probe, relevant for e.g. for UProbe links
+	ProbeLinkName string
+
 	// ExtraMeta stores extra meta info that may have been produced by a
 	// `SampleAttrProducer` instance. May be nil.
 	ExtraMeta any
 }
 
+type Origin struct {
+	Origin        libpf.Origin
+	ProbeLinkName string
+}
+
 // TraceEventsTree stores samples and their related metadata in a tree-like
 // structure optimized for the OTel Profiling protocol representation.
-type TraceEventsTree map[ContainerID]map[libpf.Origin]KeyToEventMapping
+type TraceEventsTree map[ContainerID]map[Origin]KeyToEventMapping
 
 // ContainerID represents an extracted key from /proc/<PID>/cgroup.
 type ContainerID string

support/ebpf/bpfdefs.h

@@ -129,6 +129,8 @@ static long (*bpf_probe_read_user)(void *dst, int size, const void *unsafe_ptr)
 static long (*bpf_probe_read_kernel)(void *dst, int size, const void *unsafe_ptr) = (void *)
   BPF_FUNC_probe_read_kernel;
 
+static long (*const bpf_get_attach_cookie)(void *ctx) = (void *)BPF_FUNC_get_attach_cookie;
+
   // The sizeof in bpf_trace_printk() must include \0, else no output
   // is generated. The \n is not needed on 5.8+ kernels, but definitely on
   // 5.4 kernels.