Lazy allocate memory for process

Only allocate memory for a process if a probe requires the allocation.

Move the allocation functionality to be a method of the process.Info.
This is intrinsically linked to a single process, therefore, encapsulate
the functionality as a method.

Author: MrAlias

internal/pkg/instrumentation/manager.go

@@ -80,12 +80,6 @@ func NewManager(logger *slog.Logger, otelController *opentelemetry.Controller, p
 		return nil, err
 	}
 
-	alloc, err := process.Allocate(logger, pid)
-	if err != nil {
-		return nil, err
-	}
-	m.proc.Allocation = alloc
-
 	m.logger.Info("loaded process info", "process", m.proc)
 
 	m.filterUnusedProbes()
@@ -360,7 +354,8 @@ func (m *Manager) loadProbes() error {
 	}
 	m.exe = exe
 
-	if err := m.mount(); err != nil {
+	m.logger.Debug("Mounting bpffs")
+	if err := bpffsMount(m.proc); err != nil {
 		return err
 	}
 
@@ -380,15 +375,6 @@ func (m *Manager) loadProbes() error {
 	return nil
 }
 
-func (m *Manager) mount() error {
-	if m.proc.Allocation != nil {
-		m.logger.Debug("Mounting bpffs", "allocation", m.proc.Allocation)
-	} else {
-		m.logger.Debug("Mounting bpffs")
-	}
-	return bpffsMount(m.proc)
-}
-
 func (m *Manager) cleanup() error {
 	ctx := context.Background()
 	err := m.cp.Shutdown(context.Background())

internal/pkg/instrumentation/probe/probe.go

@@ -6,6 +6,7 @@ package probe
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -572,16 +573,46 @@ func (c StructFieldConstMinVersion) InjectOption(info *process.Info) (inject.Opt
 
 // AllocationConst is a [Const] for all the allocation details that need to be
 // injected into an eBPF program.
-type AllocationConst struct{}
+type AllocationConst struct {
+	l *slog.Logger
+}
+
+// SetLogger sets the Logger for AllocationConst operations.
+func (c AllocationConst) SetLogger(l *slog.Logger) Const {
+	c.l = l
+	return c
+}
+
+func (c AllocationConst) logger() *slog.Logger {
+	l := c.l
+	if l == nil {
+		return slog.New(discardHandlerIntance)
+	}
+	return l
+}
+
+var discardHandlerIntance = discardHandler{}
+
+// Copy of slog.DiscardHandler. Remove when support for Go < 1.24 is dropped.
+type discardHandler struct{}
+
+func (dh discardHandler) Enabled(context.Context, slog.Level) bool  { return false }
+func (dh discardHandler) Handle(context.Context, slog.Record) error { return nil }
+func (dh discardHandler) WithAttrs(attrs []slog.Attr) slog.Handler  { return dh }
+func (dh discardHandler) WithGroup(name string) slog.Handler        { return dh }
 
 // InjectOption returns the appropriately configured
 // [inject.WithAllocation] if the [process.Allocation] within td
 // are not nil. An error is returned if [process.Allocation] is nil.
 func (c AllocationConst) InjectOption(info *process.Info) (inject.Option, error) {
-	if info.Allocation == nil {
+	alloc, err := info.Alloc(c.logger())
+	if err != nil {
+		return nil, err
+	}
+	if alloc == nil {
 		return nil, errors.New("no allocation details")
 	}
-	return inject.WithAllocation(*info.Allocation), nil
+	return inject.WithAllocation(*alloc), nil
 }
 
 // KeyValConst is a [Const] for a generic key-value pair.

internal/pkg/instrumentation/testutils/testutils.go

@@ -32,10 +32,6 @@ func ProbesLoad(t *testing.T, p TestProbe, libs map[string]*semver.Version) {
 
 	info := &process.Info{
 		ID: 1,
-		Allocation: &process.Allocation{
-			StartAddr: 140434497441792,
-			EndAddr:   140434497507328,
-		},
 		Modules: map[string]*semver.Version{
 			"std": testGoVersion,
 		},

internal/pkg/process/allocate.go

@@ -21,8 +21,8 @@ type Allocation struct {
 	NumCPU    uint64
 }
 
-// Allocate allocates memory for the instrumented process.
-func Allocate(logger *slog.Logger, id ID) (*Allocation, error) {
+// allocate allocates memory for the instrumented process.
+func allocate(logger *slog.Logger, id ID) (*Allocation, error) {
 	// runtime.NumCPU doesn't query any kind of hardware or OS state,
 	// but merely uses affinity APIs to count what CPUs the given go process is available to run on.
 	// Go's implementation of runtime.NumCPU (https://github.com/golang/go/blob/48d899dcdbed4534ed942f7ec2917cf86b18af22/src/runtime/os_linux.go#L97)

internal/pkg/process/info.go

@@ -7,7 +7,10 @@ import (
 	"debug/elf"
 	"errors"
 	"fmt"
+	"log/slog"
 	"runtime/debug"
+	"sync"
+	"sync/atomic"
 
 	"github.com/Masterminds/semver/v3"
 
@@ -16,11 +19,26 @@ import (
 
 // Info are the details about a target process.
 type Info struct {
-	ID         ID
-	Functions  []*binary.Func
-	GoVersion  *semver.Version
-	Modules    map[string]*semver.Version
-	Allocation *Allocation
+	ID        ID
+	Functions []*binary.Func
+	GoVersion *semver.Version
+	Modules   map[string]*semver.Version
+
+	allocOnce onceResult[*Allocation]
+}
+
+// Alloc allocates memory for the process described by Info i.
+//
+// The underlying memory allocation is only successfully performed once for the
+// instance i. Meaning, it is safe to call this multiple times. The first
+// successful result will be returned to all subsequent calls. If an error is
+// returned, subsequent calls will re-attempt to perform the allocation.
+//
+// It is safe to call this method concurrently.
+func (i *Info) Alloc(logger *slog.Logger) (*Allocation, error) {
+	return i.allocOnce.Do(func() (*Allocation, error) {
+		return allocate(logger, i.ID)
+	})
 }
 
 // NewInfo returns a new Info with information about the process identified by
@@ -118,3 +136,38 @@ func (i *Info) GetFunctionReturns(name string) ([]uint64, error) {
 
 	return nil, fmt.Errorf("could not find returns for function %s", name)
 }
+
+// onceResult is an object that will perform exactly one action if that action
+// does not error. For errors, no state is stored and subsequent attempts will
+// be tried.
+type onceResult[T any] struct {
+	done atomic.Bool
+	mu   sync.Mutex
+	val  T
+}
+
+// Do runs f only once, and only stores the result if f returns a nil error.
+// Subsequent calls to Do will return the stored value or they will re-attempt
+// to run f and store the result if an error had been returned.
+func (o *onceResult[T]) Do(f func() (T, error)) (T, error) {
+	if !o.done.Load() {
+		// Outlined complex-path to allow inlining here.
+		return o.do(f)
+	}
+	return o.val, nil
+}
+
+func (o *onceResult[T]) do(f func() (T, error)) (T, error) {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+
+	if !o.done.Load() {
+		v, err := f()
+		if err != nil {
+			return v, err
+		}
+		o.val = v
+		o.done.Store(true)
+	}
+	return o.val, nil
+}