Update Node custom labels support to use hashmap-based labels. (#91)

The previous approach, which involved resetting the current labelset
every time a piece of javascript code ran, was too slow. So we
switched to a new approach, where we _only_ run our own code when a
new async context is created or destroyed, or when label sets change.

Since we are no longer setting the current labelset, we instead store
labelsets in a hash map from async ID -> labelset (this is implemented
in the custom_labels library). Thus, at profiling time, we need to get
the current async ID from the node.js environment, then look it up in
the hashmap. This PR accomplishes the profiling side of that change.

Author: umanwizard

customlabelstest/customlabels_test.go

@@ -22,7 +22,7 @@ func TestNativeCustomLabels(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	traceCh, _ := testutils.StartTracer(ctx, t, enabledTracers, r)
+	traceCh, _ := testutils.StartTracer(ctx, t, enabledTracers, r, false)
 	// TODO - change this to `cargo build --release --bin custom-labels-example`
 	// once we have the Rust workspace from upstream.
 	cmd := exec.Command("cargo", "build", "--release",

interpreter/customlabels/customlabels.go

@@ -16,15 +16,21 @@ import (
 )
 
 const (
-	abiVersionExport = "custom_labels_abi_version"
-	tlsExport        = "custom_labels_current_set"
+	abiVersionExport    = "custom_labels_abi_version"
+	currentSetTlsExport = "custom_labels_current_set"
+	currentHmTlsExport  = "custom_labels_async_hashmap"
 )
 
-var dsoRegex = regexp.MustCompile(`.*/libcustomlabels.*\.so|.*/customlabels\.node`)
+var dsoRegex = regexp.MustCompile(`.*/libcustomlabels.*\.so`)
+var nodeRegex = regexp.MustCompile(`.*/customlabels\.node`)
 
 type data struct {
-	abiVersionElfVA libpf.Address
-	tlsAddr         libpf.Address
+	abiVersionElfVA   libpf.Address
+	currentSetTlsAddr libpf.Address
+
+	hasCurrentHm     bool
+	currentHmTlsAddr libpf.Address
+
 	isSharedLibrary bool
 }
 
@@ -35,7 +41,6 @@ func Loader(_ interpreter.EbpfHandler, info *interpreter.LoaderInfo) (interprete
 	if err != nil {
 		return nil, err
 	}
-
 	abiVersionSym, err := ef.LookupSymbol(abiVersionExport)
 	if err != nil {
 		if errors.Is(err, pfelf.ErrSymbolNotFound) {
@@ -53,31 +58,41 @@ func Loader(_ interpreter.EbpfHandler, info *interpreter.LoaderInfo) (interprete
 	// 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
+	fn := info.FileName()
+	isNativeSharedLibrary := dsoRegex.MatchString(fn)
+	isNodeExtension := (!isNativeSharedLibrary) && nodeRegex.MatchString(fn)
+	isSharedLibrary := isNativeSharedLibrary || isNodeExtension
+
+	var currentSetTlsAddr, currentHmTlsAddr libpf.Address
+	var hasCurrentHm bool
 	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]
+		currentSetTlsAddr, ok = tlsDescs[currentSetTlsExport]
 		if !ok {
 			return nil, errors.New("failed to locate TLS descriptor for custom labels")
 		}
+		if isNodeExtension {
+			currentHmTlsAddr, hasCurrentHm = tlsDescs[currentHmTlsExport]
+		}
 	} else {
-		offset, err := ef.LookupTLSSymbolOffset(tlsExport)
+		offset, err := ef.LookupTLSSymbolOffset(currentSetTlsExport)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get tls symbol offset: %w", err)
 		}
-		tlsAddr = libpf.Address(offset)
+		currentSetTlsAddr = libpf.Address(offset)
 	}
 
 	d := data{
-		abiVersionElfVA: libpf.Address(abiVersionSym.Address),
-		tlsAddr:         tlsAddr,
-		isSharedLibrary: isSharedLibrary,
+		abiVersionElfVA:   libpf.Address(abiVersionSym.Address),
+		currentSetTlsAddr: currentSetTlsAddr,
+		hasCurrentHm:      hasCurrentHm,
+		currentHmTlsAddr:  currentHmTlsAddr,
+		isSharedLibrary:   isSharedLibrary,
 	}
 	return &d, nil
 }
@@ -98,16 +113,25 @@ func (d data) Attach(ebpf interpreter.EbpfHandler, pid libpf.PID,
 			" (only 1 is supported)", abiVersion)
 	}
 
-	var tlsOffset uint64
+	var currentSetTlsOffset uint64
 	if d.isSharedLibrary {
 		// Read TLS offset from the TLS descriptor
-		tlsOffset = rm.Uint64(bias + d.tlsAddr + 8)
+		currentSetTlsOffset = rm.Uint64(bias + d.currentSetTlsAddr + 8)
 	} else {
 		// We're in the main executable: TLS offset is known statically.
-		tlsOffset = uint64(d.tlsAddr)
+		currentSetTlsOffset = uint64(d.currentSetTlsAddr)
+	}
+
+	var currentHmTlsOffset uint64
+	if d.hasCurrentHm {
+		currentHmTlsOffset = rm.Uint64(bias + d.currentHmTlsAddr + 8)
 	}
 
-	procInfo := C.NativeCustomLabelsProcInfo{tls_offset: C.u64(tlsOffset)}
+	procInfo := C.NativeCustomLabelsProcInfo{
+		current_set_tls_offset: C.u64(currentSetTlsOffset),
+		has_current_hm:         C.bool(d.hasCurrentHm),
+		current_hm_tls_offset:  C.u64(currentHmTlsOffset),
+	}
 	if err := ebpf.UpdateProcData(libpf.CustomLabels, pid, unsafe.Pointer(&procInfo)); err != nil {
 		return nil, err
 	}