Track Probe collection+closers in Manager with ProbeReference struct

Author: damemi

internal/pkg/instrumentation/manager.go

@@ -45,7 +45,7 @@ const (
 // Manager handles the management of [probe.Probe] instances.
 type Manager struct {
 	logger          *slog.Logger
-	probes          map[probe.ID]probe.Probe
+	probes          map[probe.ID]ProbeReference
 	otelController  *opentelemetry.Controller
 	cp              ConfigProvider
 	exe             *link.Executable
@@ -59,11 +59,19 @@ type Manager struct {
 	collectionOpts  *ebpf.CollectionOptions
 }
 
+// ProbeReference is used by the Manager to track an initialized reference
+// to a Probe and its related resources such as its ebpf.Collection and io.Closers.
+type ProbeReference struct {
+	probe      probe.Probe
+	collection *ebpf.Collection
+	closers    []io.Closer
+}
+
 // NewManager returns a new [Manager].
 func NewManager(logger *slog.Logger, otelController *opentelemetry.Controller, pid process.ID, cp ConfigProvider, probes ...probe.Probe) (*Manager, error) {
 	m := &Manager{
 		logger:         logger,
-		probes:         make(map[probe.ID]probe.Probe),
+		probes:         make(map[probe.ID]ProbeReference),
 		otelController: otelController,
 		cp:             cp,
 	}
@@ -132,7 +140,10 @@ func (m *Manager) registerProbe(p probe.Probe) error {
 		return err
 	}
 
-	m.probes[id] = p
+	m.probes[id] = ProbeReference{
+		probe:   p,
+		closers: make([]io.Closer, 0),
+	}
 	return nil
 }
 
@@ -146,7 +157,7 @@ func (m *Manager) filterUnusedProbes() {
 
 	for name, inst := range m.probes {
 		funcsFound := false
-		for _, s := range inst.Manifest().Symbols {
+		for _, s := range inst.probe.Manifest().Symbols {
 			if len(s.DependsOn) == 0 {
 				if _, exists := existingFuncMap[s.Symbol]; exists {
 					funcsFound = true
@@ -213,15 +224,15 @@ func (m *Manager) applyConfig(c Config) error {
 
 		if currentlyEnabled && !newEnabled {
 			m.logger.Info("Disabling probe", "id", id)
-			err = errors.Join(err, p.Close())
+			err = errors.Join(err, m.CloseProbe(p))
 			continue
 		}
 
 		if !currentlyEnabled && newEnabled {
 			m.logger.Info("Enabling probe", "id", id)
 			err = errors.Join(err, m.LoadProbe(p, id, c))
 			if err == nil {
-				m.runProbe(p)
+				m.runProbe(p.probe)
 			}
 			continue
 		}
@@ -296,7 +307,7 @@ func (m *Manager) runProbes(ctx context.Context) (context.Context, error) {
 
 	for id, p := range m.probes {
 		if isProbeEnabled(id, m.currentConfig) {
-			m.runProbe(p)
+			m.runProbe(p.probe)
 		}
 	}
 
@@ -390,29 +401,31 @@ func (m *Manager) loadProbes() error {
 	return nil
 }
 
-func (m *Manager) LoadProbe(i probe.Probe, name probe.ID, cfg Config) error {
+func (m *Manager) LoadProbe(i ProbeReference, name probe.ID, cfg Config) error {
 	m.logger.Info("loading probe", "name", name)
 
-	err := i.Init(cfg.SamplingConfig)
+	spec, err := i.probe.Spec()
 	if err != nil {
 		return errors.Join(err, m.cleanup())
 	}
 
-	spec, err := i.Spec()
+	err = m.InjectProbeConsts(i.probe, spec)
 	if err != nil {
-		return errors.Join(err, m.cleanup())
+		return err
 	}
 
-	err = m.InjectProbeConsts(i, spec)
+	c, err := utils.InitializeEBPFCollection(spec, m.collectionOpts)
 	if err != nil {
 		return err
 	}
+	i.collection = c
 
-	c, err := utils.InitializeEBPFCollection(spec, m.collectionOpts)
+	reader, err := i.probe.InitStartupConfig(c, cfg.SamplingConfig)
 	if err != nil {
-		return err
+		return errors.Join(err, m.cleanup())
 	}
-	i.SetCollection(c)
+	i.closers = append(i.closers, reader)
+
 	return nil
 }
 
@@ -437,8 +450,8 @@ func (m *Manager) InjectProbeConsts(i probe.Probe, spec *ebpf.CollectionSpec) er
 	return inject.Constants(spec, opts...)
 }
 
-func (m *Manager) loadUprobesFromProbe(i probe.Probe) error {
-	for _, up := range i.GetUprobes() {
+func (m *Manager) loadUprobesFromProbe(i ProbeReference) error {
+	for _, up := range i.probe.GetUprobes() {
 		var skip bool
 		for _, pc := range up.PackageConstraints {
 			if pc.Constraints.Check(m.proc.Modules[pc.Package]) {
@@ -448,17 +461,17 @@ func (m *Manager) loadUprobesFromProbe(i probe.Probe) error {
 			var logFn func(string, ...any)
 			switch pc.FailureMode {
 			case probe.FailureModeIgnore:
-				logFn = i.GetLogger().Debug
+				logFn = i.probe.GetLogger().Debug
 			case probe.FailureModeWarn:
-				logFn = i.GetLogger().Warn
+				logFn = i.probe.GetLogger().Warn
 			default:
 				// Unknown and FailureModeError.
 				return fmt.Errorf("uprobe %s package constraint (%s) not met, version %v", up.Sym, pc.Constraints.String(), m.proc.Modules[pc.Package])
 			}
 
 			logFn(
 				"package constraint not meet, skipping uprobe",
-				"probe", i.Manifest().ID,
+				"probe", i.probe.Manifest().ID,
 				"symbol", up.Sym,
 				"package", pc.Package,
 				"constraint", pc.Constraints.String(),
@@ -472,22 +485,23 @@ func (m *Manager) loadUprobesFromProbe(i probe.Probe) error {
 			continue
 		}
 
-		err := m.loadUprobe(up, i.GetCollection())
+		err := m.loadUprobe(up, i.collection)
 		if err != nil {
 			var logFn func(string, ...any)
 			switch up.FailureMode {
 			case probe.FailureModeIgnore:
-				logFn = i.GetLogger().Debug
+				logFn = i.probe.GetLogger().Debug
 			case probe.FailureModeWarn:
-				logFn = i.GetLogger().Warn
+				logFn = i.probe.GetLogger().Warn
 			default:
 				// Unknown and FailureModeError.
 				return err
 			}
-			logFn("failed to load uprobe", "probe", i.Manifest().ID, "symbol", up.Sym, "error", err)
+			logFn("failed to load uprobe", "probe", i.probe.Manifest().ID, "symbol", up.Sym, "error", err)
 			continue
 		}
-		_ = i.UpdateClosers(up)
+
+		i.closers = append(i.closers, up)
 	}
 	return nil
 }
@@ -555,11 +569,26 @@ func (m *Manager) mount() error {
 	return bpffsMount(m.proc)
 }
 
+func (m *Manager) CloseProbe(p ProbeReference) error {
+	if p.collection != nil {
+		p.collection.Close()
+	}
+
+	var err error
+	for _, c := range p.closers {
+		err = errors.Join(err, c.Close())
+	}
+	if err == nil {
+		p.probe.GetLogger().Debug("Closed", "Probe", p.probe.Manifest().ID)
+	}
+	return err
+}
+
 func (m *Manager) cleanup() error {
 	ctx := context.Background()
 	err := m.cp.Shutdown(context.Background())
 	for _, i := range m.probes {
-		err = errors.Join(err, i.Close())
+		err = errors.Join(err, m.CloseProbe(i))
 	}
 
 	// Wait for all probes to close so we know there is no more telemetry being

internal/pkg/instrumentation/probe/probe.go

@@ -33,36 +33,25 @@ type Probe interface {
 	// the information about the package the Probe instruments.
 	Manifest() Manifest
 
-	// Init initializes the Probe, setting up the reader, closers, and sampling config.
-	Init(*sampling.Config) error
+	// InitStartupConfig sets up initialization config options for the Probe,
+	// such as its sampling config, sets up its BPFObj as a closer, and initializes
+	// the Probe's reader, returning it as an io.Closer.
+	InitStartupConfig(*ebpf.Collection, *sampling.Config) (io.Closer, error)
 
 	// Run runs the events processing loop.
 	Run(func(ptrace.ScopeSpans))
 
-	// Close stops the Probe.
-	Close() error
-
 	// GetLogger returns the *slog.Logger associated with the Probe.
 	GetLogger() *slog.Logger
 
 	// Spec returns the *ebpf.CollectionSpec for the Probe.
 	Spec() (*ebpf.CollectionSpec, error)
 
-	// SetCollection sets the *ebpf.Collection for the Probe.
-	SetCollection(*ebpf.Collection)
-
-	// GetCollection returns the *ebpf.Collection for the Probe.
-	GetCollection() *ebpf.Collection
-
 	// GetUprobes returns a list of *Uprobes for the Probe.
 	GetUprobes() []*Uprobe
 
 	// GetConsts returns a list of Consts for the Probe.
 	GetConsts() []Const
-
-	// UpdateClosers updates the closers for the Probe to the io.Closers passed to it,
-	// and returns the new list of io.Closers for the Probe.
-	UpdateClosers(...io.Closer) []io.Closer
 }
 
 // Base is a base implementation of [Probe].
@@ -92,7 +81,6 @@ type Base[BPFObj any, BPFEvent any] struct {
 	ProcessRecord func(perf.Record) (*BPFEvent, error)
 
 	reader          *perf.Reader
-	collection      *ebpf.Collection
 	closers         []io.Closer
 	samplingManager *sampling.Manager
 }
@@ -130,40 +118,33 @@ func (i *Base[BPFObj, BPFEvent]) Spec() (*ebpf.CollectionSpec, error) {
 	return i.SpecFn()
 }
 
-// Init initializes the Probe, setting up the io.Closers, Reader, and Sampling config.
-func (i *Base[BPFObj, BPFEvent]) Init(sampler *sampling.Config) error {
+func (i *Base[BPFObj, BPFEvent]) InitStartupConfig(c *ebpf.Collection, sampler *sampling.Config) (io.Closer, error) {
 	obj := new(BPFObj)
 	if c, ok := ((interface{})(obj)).(io.Closer); ok {
 		i.closers = append(i.closers, c)
 	}
 
-	err := i.initReader()
+	samplingManager, err := sampling.NewSamplingManager(c, sampler)
 	if err != nil {
-		return err
+		return nil, err
 	}
+	i.samplingManager = samplingManager
 
-	i.samplingManager, err = sampling.NewSamplingManager(i.collection, sampler)
+	buf, ok := c.Maps[DefaultBufferMapName]
+	if !ok {
+		return nil, fmt.Errorf("%s map not found", DefaultBufferMapName)
+	}
+	i.reader, err = perf.NewReader(buf, PerfBufferDefaultSizeInPages*os.Getpagesize())
 	if err != nil {
-		return err
+		return nil, err
 	}
-
-	i.closers = append(i.closers, i.reader)
-
-	return nil
+	return i.reader, nil
 }
 
 func (i *Base[BPFObj, BPFEvent]) GetLogger() *slog.Logger {
 	return i.Logger
 }
 
-func (i *Base[BPFObj, BPFEvent]) SetCollection(c *ebpf.Collection) {
-	i.collection = c
-}
-
-func (i *Base[BPFObj, BPFEvent]) GetCollection() *ebpf.Collection {
-	return i.collection
-}
-
 func (i *Base[BPFObj, BPFEvent]) GetUprobes() []*Uprobe {
 	return i.Uprobes
 }
@@ -177,20 +158,6 @@ func (i *Base[BPFObj, BPFEvent]) GetConsts() []Const {
 	return i.Consts
 }
 
-func (i *Base[BPFObj, BPFEvent]) initReader() error {
-	buf, ok := i.collection.Maps[DefaultBufferMapName]
-	if !ok {
-		return fmt.Errorf("%s map not found", DefaultBufferMapName)
-	}
-	var err error
-	i.reader, err = perf.NewReader(buf, PerfBufferDefaultSizeInPages*os.Getpagesize())
-	if err != nil {
-		return err
-	}
-	i.closers = append(i.closers, i.reader)
-	return nil
-}
-
 // read reads a new BPFEvent from the perf Reader.
 func (i *Base[BPFObj, BPFEvent]) read() (*BPFEvent, error) {
 	record, err := i.reader.Read()
@@ -221,21 +188,6 @@ func (i *Base[BPFObj, BPFEvent]) read() (*BPFEvent, error) {
 	return event, nil
 }
 
-// Close stops the Probe.
-func (i *Base[BPFObj, BPFEvent]) Close() error {
-	if i.collection != nil {
-		i.collection.Close()
-	}
-	var err error
-	for _, c := range i.closers {
-		err = errors.Join(err, c.Close())
-	}
-	if err == nil {
-		i.Logger.Debug("Closed", "Probe", i.ID)
-	}
-	return err
-}
-
 type SpanProducer[BPFObj any, BPFEvent any] struct {
 	Base[BPFObj, BPFEvent]