usm: Better handle shared maps and probes in configureManagerWithSupp…

…ortedProtocols (#34715)

pkg/network/usm/ebpf_main.go

@@ -325,38 +325,75 @@ func (e *ebpfProgram) getProtocolsForBuildMode() ([]*protocols.ProtocolSpec, []*
 // TailCalls to the program's lists. Also, we're providing a cleanup method (the return value) which allows removal
 // of the elements we added in case of a failure in the initialization.
 func (e *ebpfProgram) configureManagerWithSupportedProtocols(protocols []*protocols.ProtocolSpec) func() {
+	// Track already existing items
+	existingMaps := make(map[string]struct{})
+	existingProbes := make(map[string]struct{})
+	existingTailCalls := make(map[string]struct{})
+
+	// Populate sets with existing elements
+	for _, m := range e.Maps {
+		existingMaps[m.Name] = struct{}{}
+	}
+	for _, p := range e.Probes {
+		existingProbes[p.EBPFFuncName] = struct{}{}
+	}
+	for _, tc := range e.tailCallRouter {
+		existingTailCalls[tc.ProbeIdentificationPair.EBPFFuncName] = struct{}{}
+	}
+
+	// Track newly added elements for cleanup
+	var addedMaps []*manager.Map
+	var addedProbes []*manager.Probe
+	var addedTailCalls []manager.TailCallRoute
+
 	for _, spec := range protocols {
-		e.Maps = append(e.Maps, spec.Maps...)
-		e.Probes = append(e.Probes, spec.Probes...)
-		e.tailCallRouter = append(e.tailCallRouter, spec.TailCalls...)
+		for _, m := range spec.Maps {
+			if _, exists := existingMaps[m.Name]; !exists {
+				e.Maps = append(e.Maps, m)
+				addedMaps = append(addedMaps, m)
+				existingMaps[m.Name] = struct{}{}
+			}
+		}
+
+		for _, p := range spec.Probes {
+			if _, exists := existingProbes[p.EBPFFuncName]; !exists {
+				e.Probes = append(e.Probes, p)
+				addedProbes = append(addedProbes, p)
+				existingProbes[p.EBPFFuncName] = struct{}{}
+			}
+		}
+
+		for _, tc := range spec.TailCalls {
+			if _, exists := existingTailCalls[tc.ProbeIdentificationPair.EBPFFuncName]; !exists {
+				e.tailCallRouter = append(e.tailCallRouter, tc)
+				addedTailCalls = append(addedTailCalls, tc)
+				existingTailCalls[tc.ProbeIdentificationPair.EBPFFuncName] = struct{}{}
+			}
+		}
 	}
+
+	// Cleanup function to remove only what was added
 	return func() {
 		e.Maps = slices.DeleteFunc(e.Maps, func(m *manager.Map) bool {
-			for _, spec := range protocols {
-				for _, specMap := range spec.Maps {
-					if m.Name == specMap.Name {
-						return true
-					}
+			for _, added := range addedMaps {
+				if m.Name == added.Name {
+					return true
 				}
 			}
 			return false
 		})
 		e.Probes = slices.DeleteFunc(e.Probes, func(p *manager.Probe) bool {
-			for _, spec := range protocols {
-				for _, probe := range spec.Probes {
-					if p.EBPFFuncName == probe.EBPFFuncName {
-						return true
-					}
+			for _, added := range addedProbes {
+				if p.EBPFFuncName == added.EBPFFuncName {
+					return true
 				}
 			}
 			return false
 		})
 		e.tailCallRouter = slices.DeleteFunc(e.tailCallRouter, func(tc manager.TailCallRoute) bool {
-			for _, spec := range protocols {
-				for _, tailCall := range spec.TailCalls {
-					if tc.ProbeIdentificationPair.EBPFFuncName == tailCall.ProbeIdentificationPair.EBPFFuncName {
-						return true
-					}
+			for _, added := range addedTailCalls {
+				if tc.ProbeIdentificationPair.EBPFFuncName == added.ProbeIdentificationPair.EBPFFuncName {
+					return true
 				}
 			}
 			return false

pkg/network/usm/ebpf_main_test.go

@@ -0,0 +1,139 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2025-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package usm
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	manager "github.com/DataDog/ebpf-manager"
+
+	"github.com/DataDog/datadog-agent/pkg/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/network/protocols"
+)
+
+func newMap(name string) *manager.Map { return &manager.Map{Name: name} }
+
+func newProbe(name string) *manager.Probe {
+	return &manager.Probe{ProbeIdentificationPair: manager.ProbeIdentificationPair{EBPFFuncName: name}}
+}
+
+func newTailCall(name string) manager.TailCallRoute {
+	return manager.TailCallRoute{ProbeIdentificationPair: manager.ProbeIdentificationPair{EBPFFuncName: name}}
+}
+
+func newEmptyEBPFProgram() *ebpfProgram {
+	return &ebpfProgram{Manager: &ebpf.Manager{Manager: &manager.Manager{}}}
+}
+
+// Common Assertions
+func assertContains(t *testing.T, e *ebpfProgram, maps, probes, calls int) {
+	require.Len(t, e.Maps, maps)
+	require.Len(t, e.Probes, probes)
+	require.Len(t, e.tailCallRouter, calls)
+}
+
+func TestConfigureManagerWithSupportedProtocols_Sanity(t *testing.T) {
+	e := newEmptyEBPFProgram()
+
+	protocolSpecs := []*protocols.ProtocolSpec{
+		{
+			Maps: []*manager.Map{
+				newMap("map1"),
+			},
+			Probes: []*manager.Probe{
+				newProbe("probe1"),
+			},
+			TailCalls: []manager.TailCallRoute{
+				newTailCall("tailcall1"),
+			},
+		},
+	}
+
+	cleanup := e.configureManagerWithSupportedProtocols(protocolSpecs)
+	assertContains(t, e, 1, 1, 1)
+	cleanup()
+	assertContains(t, e, 0, 0, 0)
+}
+
+func TestConfigureManagerWithSupportedProtocols_NoDuplicates(t *testing.T) {
+	e := newEmptyEBPFProgram()
+
+	protocolSpecs := []*protocols.ProtocolSpec{
+		{
+			Maps: []*manager.Map{
+				newMap("map1"),
+			},
+			Probes: []*manager.Probe{
+				newProbe("probe1"),
+			},
+			TailCalls: []manager.TailCallRoute{
+				newTailCall("tailcall1"),
+			},
+		},
+		{
+			Maps: []*manager.Map{
+				newMap("map1"), // Duplicate
+			},
+			Probes: []*manager.Probe{
+				newProbe("probe1"), // Duplicate
+			},
+			TailCalls: []manager.TailCallRoute{
+				newTailCall("tailcall1"), // Duplicate
+			},
+		},
+	}
+
+	cleanup := e.configureManagerWithSupportedProtocols(protocolSpecs)
+	assertContains(t, e, 1, 1, 1)
+	cleanup()
+	assertContains(t, e, 0, 0, 0)
+}
+
+func TestConfigureManagerWithSupportedProtocols_CleanupOnlyRemovesAdded(t *testing.T) {
+	e := &ebpfProgram{
+		Manager: &ebpf.Manager{
+			Manager: &manager.Manager{
+				Maps: []*manager.Map{
+					newMap("existingMap"),
+				},
+				Probes: []*manager.Probe{
+					newProbe("existingProbe"),
+				},
+			},
+		},
+		tailCallRouter: []manager.TailCallRoute{
+			newTailCall("existingTailCall"),
+		},
+	}
+
+	protocolSpecs := []*protocols.ProtocolSpec{
+		{
+			Maps: []*manager.Map{
+				newMap("newMap"),
+			},
+			Probes: []*manager.Probe{
+				newProbe("newProbe"),
+			},
+			TailCalls: []manager.TailCallRoute{
+				newTailCall("newTailCall"),
+			},
+		},
+	}
+
+	cleanup := e.configureManagerWithSupportedProtocols(protocolSpecs)
+	assertContains(t, e, 2, 2, 2)
+
+	cleanup()
+	assertContains(t, e, 1, 1, 1)
+	assert.Equal(t, "existingMap", e.Maps[0].Name)
+	assert.Equal(t, "existingProbe", e.Probes[0].EBPFFuncName)
+	assert.Equal(t, "existingTailCall", e.tailCallRouter[0].ProbeIdentificationPair.EBPFFuncName)
+}