From 837eed06eb9eaeb9e583c8f6b4c7fd7d84c5858f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillermo=20Julia=CC=81n?= <guillermo.julian@datadoghq.com>
Date: Fri, 20 Dec 2024 07:21:21 -0800
Subject: [PATCH] Add AfterStop

---
 pkg/ebpf/helper_call_patcher.go |  4 ++++
 pkg/ebpf/manager.go             | 13 ++++++++++++-
 pkg/ebpf/manager_test.go        |  4 ++++
 pkg/ebpf/printk_patcher.go      |  5 +++++
 pkg/ebpf/telemetry/modifier.go  |  5 +++++
 5 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/pkg/ebpf/helper_call_patcher.go b/pkg/ebpf/helper_call_patcher.go
index 09325f2d72580..3ed737c592045 100644
--- a/pkg/ebpf/helper_call_patcher.go
+++ b/pkg/ebpf/helper_call_patcher.go
@@ -92,3 +92,7 @@ func (h *helperCallRemover) String() string {
 func (h *helperCallRemover) BeforeStop(*manager.Manager, names.ModuleName) error {
 	return nil
 }
+
+func (h *helperCallRemover) AfterStop(*manager.Manager, names.ModuleName) error {
+	return nil
+}
diff --git a/pkg/ebpf/manager.go b/pkg/ebpf/manager.go
index 77f005a660239..a280004d7e429 100644
--- a/pkg/ebpf/manager.go
+++ b/pkg/ebpf/manager.go
@@ -75,6 +75,10 @@ type Modifier interface {
 	// BeforeStop is called after the ebpf.Manager.Stop call. Errors on these modifiers
 	// will not prevent the manager from stopping, but they will be reported.
 	BeforeStop(*manager.Manager, names.ModuleName) error
+
+	// AfterStop is called after the ebpf.Manager.Stop call. If the manager returns
+	// an error on stop, this function will still be called.
+	AfterStop(*manager.Manager, names.ModuleName) error
 }
 
 // InitWithOptions is a wrapper around ebpf-manager.Manager.InitWithOptions
@@ -111,7 +115,7 @@ func (m *Manager) Stop(clean manager.MapCleanupType) error {
 	var errs error
 
 	for _, mod := range m.EnabledModifiers {
-		log.Warnf("Running %s manager modifier BeforeStop", mod)
+		log.Tracef("Running %s manager modifier BeforeStop", mod)
 		if err := mod.BeforeStop(m.Manager, m.Name); err != nil {
 			errs = errors.Join(errs, fmt.Errorf("error running %s manager modifier BeforeStop(): %s", mod, err))
 		}
@@ -121,5 +125,12 @@ func (m *Manager) Stop(clean manager.MapCleanupType) error {
 		errs = errors.Join(errs, fmt.Errorf("failed to stop manager %w", err))
 	}
 
+	for _, mod := range m.EnabledModifiers {
+		log.Tracef("Running %s manager modifier AfterStop", mod)
+		if err := mod.AfterStop(m.Manager, m.Name); err != nil {
+			errs = errors.Join(errs, fmt.Errorf("error running %s manager modifier AfterStop(): %s", mod, err))
+		}
+	}
+
 	return errs
 }
diff --git a/pkg/ebpf/manager_test.go b/pkg/ebpf/manager_test.go
index aaf5b931bde14..6a4bdc8ebaf8a 100644
--- a/pkg/ebpf/manager_test.go
+++ b/pkg/ebpf/manager_test.go
@@ -40,6 +40,10 @@ func (t *dummyModifier) BeforeStop(_ *manager.Manager, _ names.ModuleName) error
 	return nil
 }
 
+func (t *dummyModifier) AfterStop(_ *manager.Manager, _ names.ModuleName) error {
+	return nil
+}
+
 func TestNewManagerWithDefault(t *testing.T) {
 	type args struct {
 		mgr       *manager.Manager
diff --git a/pkg/ebpf/printk_patcher.go b/pkg/ebpf/printk_patcher.go
index 451e0835e775d..67f1c8df74f02 100644
--- a/pkg/ebpf/printk_patcher.go
+++ b/pkg/ebpf/printk_patcher.go
@@ -232,3 +232,8 @@ func (t *PrintkPatcherModifier) AfterInit(_ *manager.Manager, _ names.ModuleName
 func (t *PrintkPatcherModifier) BeforeStop(_ *manager.Manager, _ names.ModuleName) error {
 	return nil
 }
+
+// AfterStop is a no-op for this modifier
+func (t *PrintkPatcherModifier) AfterStop(_ *manager.Manager, _ names.ModuleName) error {
+	return nil
+}
diff --git a/pkg/ebpf/telemetry/modifier.go b/pkg/ebpf/telemetry/modifier.go
index b2fa9815caecd..465bb9d9b29b5 100644
--- a/pkg/ebpf/telemetry/modifier.go
+++ b/pkg/ebpf/telemetry/modifier.go
@@ -140,3 +140,8 @@ func (t *ErrorsTelemetryModifier) AfterInit(m *manager.Manager, module names.Mod
 func (t *ErrorsTelemetryModifier) BeforeStop(_ *manager.Manager, _ names.ModuleName) error {
 	return nil
 }
+
+// AfterStop is a no-op
+func (t *ErrorsTelemetryModifier) AfterStop(_ *manager.Manager, _ names.ModuleName) error {
+	return nil
+}