feat: apply custom labels patch

.github/workflows/unit-test-on-pull-request.yml

@@ -58,6 +58,8 @@ jobs:
         uses: ./.github/workflows/env
       - name: Tests
         run: make test TARGET_ARCH=${{ matrix.target_arch }}
+      - name: sudo tests
+        run: make sudo-tests
 
   build-integration-test-binaries:
     name: Build integration test binaries (${{ matrix.target_arch }})

Makefile

@@ -76,6 +76,16 @@ linter-version:
 test: generate ebpf test-deps
 	go test $(GO_FLAGS) ./...
 
+SUDOTEST_DIRS := ./customlabelstest
+
+target/release/custom-labels-example:
+	cargo build --release --bin custom-labels-example
+
+sudo-tests: generate ebpf target/release/custom-labels-example
+	$(foreach test_dir, $(SUDOTEST_DIRS), \
+		sudo go test $(GO_FLAGS) -tags $(GO_TAGS) "$(test_dir)" \
+	)
+
 TESTDATA_DIRS:= \
 	nativeunwind/elfunwindinfo/testdata \
 	libpf/pfelf/testdata \

customlabelstest/customlabels_test.go

@@ -0,0 +1,80 @@
+package customlabelstest
+
+import (
+	"context"
+	"os/exec"
+	"regexp"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/open-telemetry/opentelemetry-ebpf-profiler/testutils"
+	tracertypes "github.com/open-telemetry/opentelemetry-ebpf-profiler/tracer/types"
+)
+
+func TestNativeCustomLabels(t *testing.T) {
+	if !testutils.IsRoot() {
+		t.Skip("root privileges required")
+	}
+
+	r := &testutils.MockReporter{}
+	enabledTracers, _ := tracertypes.Parse("all")
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	traceCh, _ := testutils.StartTracer(ctx, t, enabledTracers, r)
+
+	errCh := make(chan error, 1)
+
+	cmd := exec.CommandContext(ctx, "../target/release/custom-labels-example")
+	err := cmd.Start()
+	require.NoError(t, err, "run 'cargo build --release --bin custom-labels-example' first")
+
+	go func() {
+		err := cmd.Wait()
+		errCh <- err
+	}()
+
+	stopCh := time.After(10 * time.Second)
+
+	re := regexp.MustCompile(`^[a-zA-Z0-9]{16}$`)
+	good := false
+Loop:
+	for {
+		select {
+		case trace, ok := <-traceCh:
+			if !ok {
+				break Loop
+			}
+			if trace == nil {
+				continue
+			}
+			if len(trace.CustomLabels) > 0 {
+				var gotL1, gotL2 bool
+				for k, v := range trace.CustomLabels {
+					switch k {
+					case "l1":
+						gotL1 = true
+						require.True(t, re.MatchString(v))
+						t.Logf("got l1, value is %s", v)
+					case "l2":
+						gotL2 = true
+						require.True(t, re.MatchString(v))
+						t.Logf("got l2, value is %s", v)
+					default:
+						require.Failf(t, "fail", "got unexpected label: %s=%s", k, v)
+					}
+				}
+				if gotL1 && gotL2 {
+					good = true
+					break Loop
+				}
+			}
+		case err := <-errCh:
+			require.Failf(t, "fail", "Failed to run custom-labels-example, err = %v", err)
+		case <-stopCh:
+			require.Fail(t, "fail", "Failed to get labels after ten seconds")
+		}
+	}
+	require.True(t, good)
+}

host/host.go

@@ -55,4 +55,5 @@ type Trace struct {
 	TID              libpf.PID
 	APMTraceID       libpf.APMTraceID
 	APMTransactionID libpf.APMTransactionID
+	CustomLabels     map[string]string
 }

interpreter/customlabels/customlabels.go

@@ -0,0 +1,156 @@
+package customlabels
+
+// #include <stdlib.h>
+// #include "../../support/ebpf/types.h"
+import "C"
+import (
+	"debug/elf"
+	"errors"
+	"fmt"
+	"regexp"
+	"unsafe"
+
+	"github.com/open-telemetry/opentelemetry-ebpf-profiler/interpreter"
+	"github.com/open-telemetry/opentelemetry-ebpf-profiler/libpf"
+	"github.com/open-telemetry/opentelemetry-ebpf-profiler/libpf/pfelf"
+	"github.com/open-telemetry/opentelemetry-ebpf-profiler/remotememory"
+)
+
+const (
+	abiVersionExport = "custom_labels_abi_version"
+	tlsExport        = "custom_labels_current_set"
+)
+
+var dsoRegex = regexp.MustCompile(`.*/libcustomlabels.*\.so|.*/customlabels\.node`)
+
+type data struct {
+	abiVersionElfVA libpf.Address
+	tlsAddr         libpf.Address
+	isSharedLibrary bool
+}
+
+var _ interpreter.Data = &data{}
+
+func roundUp(multiple, value uint64) uint64 {
+	if multiple == 0 {
+		return value
+	}
+	return (value + multiple - 1) / multiple * multiple
+}
+
+func Loader(_ interpreter.EbpfHandler, info *interpreter.LoaderInfo) (interpreter.Data, error) {
+	ef, err := info.GetELF()
+	if err != nil {
+		return nil, err
+	}
+
+	abiVersionSym, err := ef.LookupSymbol(abiVersionExport)
+	if err != nil {
+		if errors.Is(err, pfelf.ErrSymbolNotFound) {
+			return nil, nil
+		}
+
+		return nil, err
+	}
+
+	if abiVersionSym.Size != 4 {
+		return nil, fmt.Errorf("abi version export has wrong size %d", abiVersionSym.Size)
+	}
+
+	// If this is the libcustomlabels.so library, we are using
+	// global-dynamic TLS model and have to look up the TLS descriptor.
+	// Otherwise, assume we're the main binary and just look up the
+	// symbol.
+	isSharedLibrary := dsoRegex.MatchString(info.FileName())
+	var tlsAddr libpf.Address
+	if isSharedLibrary {
+		// Resolve thread info TLS export.
+		tlsDescs, err := ef.TLSDescriptors()
+		if err != nil {
+			return nil, errors.New("failed to extract TLS descriptors")
+		}
+		var ok bool
+		tlsAddr, ok = tlsDescs[tlsExport]
+		if !ok {
+			return nil, errors.New("failed to locate TLS descriptor for custom labels")
+		}
+	} else {
+		tlsSym, err := ef.LookupSymbol(tlsExport)
+		if err != nil {
+			return nil, err
+		}
+		if ef.Machine == elf.EM_AARCH64 {
+			tlsAddr = libpf.Address(tlsSym.Address + 16)
+		} else if ef.Machine == elf.EM_X86_64 {
+			// Symbol addresses are relative to the start of the
+			// thread-local storage image, but the thread pointer points to the _end_
+			// of the image. So we need to find the size of the image in order to know where the
+			// beginning is.
+			//
+			// The image is just .tdata followed by .tbss,
+			// but we also have to respect the alignment.
+			tbss, err := ef.Tbss()
+			if err != nil {
+				return nil, err
+			}
+			tdata, err := ef.Tdata()
+			var tdataSize uint64
+			if err != nil {
+				// No Tdata is ok, it's the same as size 0
+				if err != pfelf.ErrNoTdata {
+					return nil, err
+				}
+			} else {
+				tdataSize = tdata.Size
+			}
+			imageSize := roundUp(tbss.Addralign, tdataSize) + tbss.Size
+			tlsAddr = libpf.Address(int64(tlsSym.Address) - int64(imageSize))
+		} else {
+			return nil, fmt.Errorf("unrecognized machine: %s", ef.Machine.String())
+		}
+	}
+
+	d := data{
+		abiVersionElfVA: libpf.Address(abiVersionSym.Address),
+		tlsAddr:         tlsAddr,
+		isSharedLibrary: isSharedLibrary,
+	}
+	return &d, nil
+}
+
+type instance struct {
+	interpreter.InstanceStubs
+}
+
+func (d data) Attach(ebpf interpreter.EbpfHandler, pid libpf.PID,
+	bias libpf.Address, rm remotememory.RemoteMemory) (interpreter.Instance, error) {
+	abiVersion, err := rm.Uint32Checked(bias + d.abiVersionElfVA)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read custom labels ABI version: %w", err)
+	}
+
+	if abiVersion != 1 {
+		return nil, fmt.Errorf("unsupported custom labels ABI version: %d"+
+			" (only 1 is supported)", abiVersion)
+	}
+
+	var tlsOffset uint64
+	if d.isSharedLibrary {
+		// Read TLS offset from the TLS descriptor
+		tlsOffset = rm.Uint64(bias + d.tlsAddr + 8)
+	} else {
+		// We're in the main executable: TLS offset is known statically.
+		tlsOffset = uint64(d.tlsAddr)
+	}
+
+	procInfo := C.NativeCustomLabelsProcInfo{tls_offset: C.u64(tlsOffset)}
+	if err := ebpf.UpdateProcData(libpf.CustomLabels, pid, unsafe.Pointer(&procInfo)); err != nil {
+		return nil, err
+	}
+
+	return &instance{}, nil
+}
+
+func (i *instance) Detach(ebpf interpreter.EbpfHandler, pid libpf.PID) error {
+	return ebpf.DeleteProcData(libpf.CustomLabels, pid)
+}

libpf/interpretertype.go

@@ -40,6 +40,9 @@ const (
 
 	// APMInt identifies the pseudo-interpreter for the APM integration.
 	APMInt InterpreterType = 0x100
+
+	// CustomLabels identifies the pseudo-interpreter for native custom labels support.
+	CustomLabels InterpreterType = 0x102
 )
 
 // Frame converts the interpreter type into the corresponding frame type.

libpf/pfelf/file.go

@@ -54,6 +54,12 @@ var ErrSymbolNotFound = errors.New("symbol not found")
 // ErrNotELF is returned when the file is not an ELF
 var ErrNotELF = errors.New("not an ELF file")
 
+// ErrNoTbss is returned when the tbss section cannot be found
+var ErrNoTbss = errors.New("no thread-local uninitialized data section (tbss)")
+
+// ErrNoTdata is returned when the tdata section cannot be found
+var ErrNoTdata = errors.New("no thread-local initialized data section (tdata)")
+
 // File represents an open ELF file
 type File struct {
 	// closer is called internally when resources for this File are to be released
@@ -413,6 +419,32 @@ func (f *File) Section(name string) *Section {
 	return nil
 }
 
+// Tbss gets the thread-local uninitialized data section
+func (f *File) Tbss() (*Section, error) {
+	if err := f.LoadSections(); err != nil {
+		return nil, err
+	}
+	for _, sec := range f.Sections {
+		if sec.Type == elf.SHT_NOBITS && sec.Flags&elf.SHF_TLS != 0 {
+			return &sec, nil
+		}
+	}
+	return nil, ErrNoTbss
+}
+
+// Tdata gets the thread-local initialized data section
+func (f *File) Tdata() (*Section, error) {
+	if err := f.LoadSections(); err != nil {
+		return nil, err
+	}
+	for _, sec := range f.Sections {
+		if sec.Type == elf.SHT_PROGBITS && sec.Flags&elf.SHF_TLS != 0 {
+			return &sec, nil
+		}
+	}
+	return nil, ErrNoTdata
+}
+
 // ReadVirtualMemory reads bytes from given virtual address
 func (f *File) ReadVirtualMemory(p []byte, addr int64) (int, error) {
 	if len(p) == 0 {

libpf/trace.go

@@ -14,6 +14,7 @@ type Trace struct {
 	MappingEnd         []Address
 	MappingFileOffsets []uint64
 	Hash               TraceHash
+	CustomLabels       map[string]string
 }
 
 // AppendFrame appends a frame to the columnar frame array without mapping information.

processmanager/ebpf/ebpf.go

@@ -103,6 +103,7 @@ type ebpfMapsImpl struct {
 	rubyProcs          *cebpf.Map
 	v8Procs            *cebpf.Map
 	apmIntProcs        *cebpf.Map
+	clProcs            *cebpf.Map
 
 	// Stackdelta and process related eBPF maps
 	exeIDToStackDeltaMaps []*cebpf.Map
@@ -203,6 +204,12 @@ func LoadMaps(ctx context.Context, maps map[string]*cebpf.Map) (EbpfHandler, err
 	}
 	impl.apmIntProcs = apmIntProcs
 
+	clProcs, ok := maps["cl_procs"]
+	if !ok {
+		log.Fatalf("Map cl_procs is not available")
+	}
+	impl.clProcs = clProcs
+
 	impl.stackDeltaPageToInfo, ok = maps["stack_delta_page_to_info"]
 	if !ok {
 		log.Fatalf("Map stack_delta_page_to_info is not available")
@@ -294,6 +301,8 @@ func (impl *ebpfMapsImpl) getInterpreterTypeMap(typ libpf.InterpreterType) (*ceb
 		return impl.v8Procs, nil
 	case libpf.APMInt:
 		return impl.apmIntProcs, nil
+	case libpf.CustomLabels:
+		return impl.clProcs, nil
 	default:
 		return nil, fmt.Errorf("type %d is not (yet) supported", typ)
 	}

processmanager/execinfomanager/manager.go

@@ -18,6 +18,7 @@ import (
 	"github.com/open-telemetry/opentelemetry-ebpf-profiler/host"
 	"github.com/open-telemetry/opentelemetry-ebpf-profiler/interpreter"
 	"github.com/open-telemetry/opentelemetry-ebpf-profiler/interpreter/apmint"
+	"github.com/open-telemetry/opentelemetry-ebpf-profiler/interpreter/customlabels"
 	"github.com/open-telemetry/opentelemetry-ebpf-profiler/interpreter/dotnet"
 	"github.com/open-telemetry/opentelemetry-ebpf-profiler/interpreter/hotspot"
 	"github.com/open-telemetry/opentelemetry-ebpf-profiler/interpreter/nodev8"
@@ -100,6 +101,7 @@ func NewExecutableInfoManager(
 	sdp nativeunwind.StackDeltaProvider,
 	ebpf pmebpf.EbpfHandler,
 	includeTracers types.IncludedTracers,
+	collectCustomLabels bool,
 ) (*ExecutableInfoManager, error) {
 	// Initialize interpreter loaders.
 	interpreterLoaders := make([]interpreter.Loader, 0)
@@ -126,6 +128,9 @@ func NewExecutableInfoManager(
 	}
 
 	interpreterLoaders = append(interpreterLoaders, apmint.Loader)
+	if collectCustomLabels {
+		interpreterLoaders = append(interpreterLoaders, customlabels.Loader)
+	}
 
 	deferredFileIDs, err := lru.NewSynced[host.FileID, libpf.Void](deferredFileIDSize,
 		func(id host.FileID) uint32 { return uint32(id) })