[EBPF] Add BeforeStop/AfterStop method to modifiers (DataDog#32453)

pkg/ebpf/helper_call_patcher.go

@@ -40,7 +40,7 @@ var replaceIns = asm.Mov.Imm(asm.R0, 0)
 // conditionally select eBPF helpers. This should be regarded as a last resort
 // when the aforementioned options don't apply (prebuilt artifacts, for
 // example).
-func NewHelperCallRemover(helpers ...asm.BuiltinFunc) Modifier {
+func NewHelperCallRemover(helpers ...asm.BuiltinFunc) ModifierBeforeInit {
 	return &helperCallRemover{
 		helpers: helpers,
 	}
@@ -81,10 +81,6 @@ func (h *helperCallRemover) BeforeInit(m *manager.Manager, _ names.ModuleName, _
 	return nil
 }
 
-func (h *helperCallRemover) AfterInit(*manager.Manager, names.ModuleName, *manager.Options) error {
-	return nil
-}
-
 func (h *helperCallRemover) String() string {
 	return fmt.Sprintf("HelperCallRemover[%+v]", h.helpers)
 }

pkg/ebpf/manager.go

@@ -8,6 +8,7 @@
 package ebpf
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"sync"
@@ -51,27 +52,76 @@ func NewManagerWithDefault(mgr *manager.Manager, name string, modifiers ...Modif
 	return NewManager(mgr, name, append(defaultModifiers, modifiers...)...)
 }
 
-// Modifier is an interface that can be implemented by a package to
-// add functionality to the ebpf.Manager. It exposes a name to identify the modifier,
-// two functions that will be called before and after the ebpf.Manager.InitWithOptions
-// call, and a function that will be called when the manager is stopped.
-// Note regarding internal state of the modifier: if the modifier is added to the list of modifiers
-// enabled by default (pkg/ebpf/ebpf.go:registerDefaultModifiers), all managers with those default modifiers
-// will share the same instance of the modifier. On the other hand, if the modifier is added to a specific
-// manager, it can have its own instance of the modifier, unless the caller explicitly uses the same modifier
-// instance with different managers. In other words, if the modifier is to have any internal state specific to
-// each manager, it should not be added to the list of default modifiers, and developers using it
-// should be aware of this behavior.
+// Modifier is an interface that can be implemented by a package to add
+// functionality to the ebpf.Manager. It exposes a name to identify the
+// modifier, and then any of the functions Before/AfterInit, Before/AfterStart,
+// Before/AfterStop, that will be called at the corresponding stage of the
+// manager lifecycle. To avoid code churn and implementing unnecessary
+// functions, the Modifier interface is split into sub-interfaces, each with a
+// single function. This way, the developer can implement only the functions
+// they need, and the manager will call them at the right time. Note regarding
+// internal state of the modifier: if the modifier is added to the list of
+// modifiers enabled by default (see NewManagerWithDefault above), all managers
+// with those default modifiers will share the same instance of the modifier. On
+// the other hand, if the modifier is added to a specific manager, it can have
+// its own instance of the modifier, unless the caller explicitly uses the same
+// modifier instance with different managers. In other words, if the modifier is
+// to have any internal state specific to each manager, it should not be added
+// to the list of default modifiers, and developers using it should be aware of
+// this behavior.
 type Modifier interface {
 	fmt.Stringer
+}
+
+// ModifierBeforeInit is a sub-interface of Modifier that exposes a BeforeInit method
+type ModifierBeforeInit interface {
+	Modifier
+
 	// BeforeInit is called before the ebpf.Manager.InitWithOptions call
-	// names.ModuleName refers to the name associated with Manager instance.
+	// names.ModuleName refers to the name associated with Manager instance. An
+	// error returned from this function will stop the initialization process.
 	BeforeInit(*manager.Manager, names.ModuleName, *manager.Options) error
+}
+
+// ModifierAfterInit is a sub-interface of Modifier that exposes an AfterInit method
+type ModifierAfterInit interface {
+	Modifier
 
 	// AfterInit is called after the ebpf.Manager.InitWithOptions call
 	AfterInit(*manager.Manager, names.ModuleName, *manager.Options) error
 }
 
+// ModifierBeforeStop is a sub-interface of Modifier that exposes a BeforeStop method
+type ModifierBeforeStop interface {
+	Modifier
+
+	// BeforeStop is called before the ebpf.Manager.Stop call. An error returned
+	// from this function will not prevent the manager from stopping, but it will
+	// be logged.
+	BeforeStop(*manager.Manager, names.ModuleName, manager.MapCleanupType) error
+}
+
+// ModifierAfterStop is a sub-interface of Modifier that exposes an AfterStop method
+type ModifierAfterStop interface {
+	Modifier
+
+	// AfterStop is called after the ebpf.Manager.Stop call. An error returned
+	// from this function will be logged.
+	AfterStop(*manager.Manager, names.ModuleName, manager.MapCleanupType) error
+}
+
+func runModifiersOfType[K Modifier](modifiers []Modifier, funcName string, runner func(K) error) error {
+	var errs error
+	for _, mod := range modifiers {
+		if as, ok := mod.(K); ok {
+			if err := runner(as); err != nil {
+				errs = errors.Join(errs, fmt.Errorf("error running %s manager modifier %s: %w", mod, funcName, err))
+			}
+		}
+	}
+	return errs
+}
+
 // InitWithOptions is a wrapper around ebpf-manager.Manager.InitWithOptions
 func (m *Manager) InitWithOptions(bytecode io.ReaderAt, opts *manager.Options) error {
 	// we must load the ELF file before initialization,
@@ -81,22 +131,40 @@ func (m *Manager) InitWithOptions(bytecode io.ReaderAt, opts *manager.Options) e
 		return fmt.Errorf("failed to load elf from reader: %w", err)
 	}
 
-	for _, mod := range m.EnabledModifiers {
-		log.Tracef("Running %s manager modifier BeforeInit", mod)
-		if err := mod.BeforeInit(m.Manager, m.Name, opts); err != nil {
-			return fmt.Errorf("error running %s manager modifier: %w", mod, err)
-		}
+	err := runModifiersOfType(m.EnabledModifiers, "BeforeInit", func(mod ModifierBeforeInit) error {
+		return mod.BeforeInit(m.Manager, m.Name, opts)
+	})
+	if err != nil {
+		return err
 	}
 
 	if err := m.Manager.InitWithOptions(nil, *opts); err != nil {
 		return err
 	}
 
-	for _, mod := range m.EnabledModifiers {
-		log.Tracef("Running %s manager modifier AfterInit", mod)
-		if err := mod.AfterInit(m.Manager, m.Name, opts); err != nil {
-			return fmt.Errorf("error running %s manager modifier: %w", mod, err)
-		}
+	return runModifiersOfType(m.EnabledModifiers, "AfterInit", func(mod ModifierAfterInit) error {
+		return mod.AfterInit(m.Manager, m.Name, opts)
+	})
+}
+
+// Stop is a wrapper around ebpf-manager.Manager.Stop
+func (m *Manager) Stop(cleanupType manager.MapCleanupType) error {
+	var errs error
+
+	err := runModifiersOfType(m.EnabledModifiers, "BeforeStop", func(mod ModifierBeforeStop) error {
+		return mod.BeforeStop(m.Manager, m.Name, cleanupType)
+	})
+	if err != nil {
+		errs = errors.Join(errs, err)
+	}
+
+	if err := m.Manager.Stop(cleanupType); err != nil {
+		errs = errors.Join(errs, fmt.Errorf("failed to stop manager %w", err))
 	}
-	return nil
+
+	err = runModifiersOfType(m.EnabledModifiers, "AfterStop", func(mod ModifierAfterStop) error {
+		return mod.AfterStop(m.Manager, m.Name, cleanupType)
+	})
+
+	return errors.Join(errs, err)
 }

pkg/ebpf/manager_test.go

@@ -9,29 +9,44 @@ package ebpf
 import (
 	"testing"
 
-	"github.com/DataDog/datadog-agent/pkg/ebpf/names"
 	manager "github.com/DataDog/ebpf-manager"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/DataDog/datadog-agent/pkg/ebpf/bytecode"
+	"github.com/DataDog/datadog-agent/pkg/ebpf/names"
 )
 
-// PrintkPatcherModifier adds an InstructionPatcher to the manager that removes the newline character from log_debug calls if needed
 type dummyModifier struct {
+	mock.Mock
 }
 
 const dummyModifierName = "DummyModifier"
 
 func (t *dummyModifier) String() string {
+	// Do not mock this method for simplicity, to avoid having to define it always
 	return dummyModifierName
 }
 
-// BeforeInit adds the patchPrintkNewline function to the manager
-func (t *dummyModifier) BeforeInit(_ *manager.Manager, _ names.ModuleName, _ *manager.Options) error {
-	return nil
+func (t *dummyModifier) BeforeInit(m *manager.Manager, name names.ModuleName, opts *manager.Options) error {
+	args := t.Called(m, name, opts)
+	return args.Error(0)
+}
+
+func (t *dummyModifier) AfterInit(m *manager.Manager, name names.ModuleName, opts *manager.Options) error {
+	args := t.Called(m, name, opts)
+	return args.Error(0)
 }
 
-// AfterInit is a no-op for this modifier
-func (t *dummyModifier) AfterInit(_ *manager.Manager, _ names.ModuleName, _ *manager.Options) error {
-	return nil
+func (t *dummyModifier) BeforeStop(m *manager.Manager, name names.ModuleName, cleanupType manager.MapCleanupType) error {
+	args := t.Called(m, name, cleanupType)
+	return args.Error(0)
+}
+
+func (t *dummyModifier) AfterStop(m *manager.Manager, name names.ModuleName, cleanupType manager.MapCleanupType) error {
+	args := t.Called(m, name, cleanupType)
+	return args.Error(0)
 }
 
 func TestNewManagerWithDefault(t *testing.T) {
@@ -69,3 +84,40 @@ func TestNewManagerWithDefault(t *testing.T) {
 		})
 	}
 }
+
+func TestManagerInitWithOptions(t *testing.T) {
+	modifier := &dummyModifier{}
+	modifier.On("BeforeInit", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	modifier.On("AfterInit", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+
+	mgr := NewManager(&manager.Manager{}, "test", modifier)
+	require.NotNil(t, mgr)
+
+	// Load a simple eBPF program to test the modifiers
+	cfg := NewConfig()
+	require.NotNil(t, cfg)
+
+	buf, err := bytecode.GetReader(cfg.BPFDir, "logdebug-test.o")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = buf.Close })
+
+	opts := manager.Options{}
+	err = mgr.InitWithOptions(buf, &opts)
+	require.NoError(t, err)
+
+	modifier.AssertExpectations(t)
+}
+
+func TestManagerStop(t *testing.T) {
+	modifier := &dummyModifier{}
+	modifier.On("BeforeStop", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	modifier.On("AfterStop", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+
+	mgr := NewManager(&manager.Manager{}, "test", modifier)
+	require.NotNil(t, mgr)
+
+	// The Stop call will fail because the manager is not initialized, but the modifiers should still be called
+	_ = mgr.Stop(manager.CleanAll)
+
+	modifier.AssertExpectations(t)
+}

pkg/ebpf/printk_patcher.go

@@ -213,6 +213,9 @@ func patchPrintkInstructions(p *ebpf.ProgramSpec) (int, error) {
 type PrintkPatcherModifier struct {
 }
 
+// ensure PrintkPatcherModifier implements the ModifierBeforeInit interface
+var _ ModifierBeforeInit = &PrintkPatcherModifier{}
+
 func (t *PrintkPatcherModifier) String() string {
 	return "PrintkPatcherModifier"
 }
@@ -222,8 +225,3 @@ func (t *PrintkPatcherModifier) BeforeInit(m *manager.Manager, _ names.ModuleNam
 	m.InstructionPatchers = append(m.InstructionPatchers, patchPrintkNewline)
 	return nil
 }
-
-// AfterInit is a no-op for this modifier
-func (t *PrintkPatcherModifier) AfterInit(_ *manager.Manager, _ names.ModuleName, _ *manager.Options) error {
-	return nil
-}