From ce240d122afa17d7bf8ac25fc47815670c816220 Mon Sep 17 00:00:00 2001 From: Guy Arbitman Date: Tue, 4 Mar 2025 17:43:29 +0200 Subject: [PATCH] usm: Better handle shared maps and probes in configureManagerWithSupportedProtocols (#34715) --- pkg/network/usm/ebpf_main.go | 73 ++++++++++++---- pkg/network/usm/ebpf_main_test.go | 139 ++++++++++++++++++++++++++++++ 2 files changed, 194 insertions(+), 18 deletions(-) create mode 100644 pkg/network/usm/ebpf_main_test.go diff --git a/pkg/network/usm/ebpf_main.go b/pkg/network/usm/ebpf_main.go index 2cfe23e271bbf6..9ab5d77edc5c40 100644 --- a/pkg/network/usm/ebpf_main.go +++ b/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 diff --git a/pkg/network/usm/ebpf_main_test.go b/pkg/network/usm/ebpf_main_test.go new file mode 100644 index 00000000000000..f49174512d231e --- /dev/null +++ b/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) +}