From ef1f063372e9718bcba54a683741a420a7425fd6 Mon Sep 17 00:00:00 2001
From: Jonathan Ribas <jonathan.ribas@datadoghq.com>
Date: Wed, 16 Oct 2024 17:42:24 +0200
Subject: [PATCH] [CWS] Fix a load controller issue where tags were resolved
 multiple times (#30178)

---
 pkg/security/security_profile/dump/activity_dump.go   | 3 ++-
 pkg/security/security_profile/dump/load_controller.go | 3 ++-
 pkg/security/security_profile/dump/manager.go         | 6 +++---
 pkg/security/tests/fake_tags_resolver.go              | 1 -
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/pkg/security/security_profile/dump/activity_dump.go b/pkg/security/security_profile/dump/activity_dump.go
index 0f345ddf77d31..b9b0f34e00e71 100644
--- a/pkg/security/security_profile/dump/activity_dump.go
+++ b/pkg/security/security_profile/dump/activity_dump.go
@@ -615,7 +615,8 @@ func (ad *ActivityDump) ResolveTags() error {
 
 // resolveTags thread unsafe version ot ResolveTags
 func (ad *ActivityDump) resolveTags() error {
-	if len(ad.Tags) >= 10 || len(ad.Metadata.ContainerID) == 0 {
+	selector := ad.GetWorkloadSelector()
+	if selector != nil {
 		return nil
 	}
 
diff --git a/pkg/security/security_profile/dump/load_controller.go b/pkg/security/security_profile/dump/load_controller.go
index fcce3b6de6a51..3d0148f168232 100644
--- a/pkg/security/security_profile/dump/load_controller.go
+++ b/pkg/security/security_profile/dump/load_controller.go
@@ -87,6 +87,7 @@ func (lc *ActivityDumpLoadController) NextPartialDump(ad *ActivityDump) *Activit
 	newDump.Metadata.ContainerID = ad.Metadata.ContainerID
 	newDump.Metadata.DifferentiateArgs = ad.Metadata.DifferentiateArgs
 	newDump.Tags = ad.Tags
+	newDump.selector = ad.selector
 
 	// copy storage requests
 	for _, reqList := range ad.StorageRequests {
@@ -101,7 +102,7 @@ func (lc *ActivityDumpLoadController) NextPartialDump(ad *ActivityDump) *Activit
 	}
 
 	// compute the duration it took to reach the dump size threshold
-	timeToThreshold := ad.End.Sub(ad.Start)
+	timeToThreshold := time.Since(ad.Start)
 
 	// set new load parameters
 	newDump.SetTimeout(ad.LoadConfig.Timeout - timeToThreshold)
diff --git a/pkg/security/security_profile/dump/manager.go b/pkg/security/security_profile/dump/manager.go
index 10af2f9d787f7..a0b8c490b965e 100644
--- a/pkg/security/security_profile/dump/manager.go
+++ b/pkg/security/security_profile/dump/manager.go
@@ -852,6 +852,9 @@ func (adm *ActivityDumpManager) triggerLoadController() {
 
 	// handle overweight dumps
 	for _, ad := range dumps {
+		// restart a new dump for the same workload
+		newDump := adm.loadController.NextPartialDump(ad)
+
 		// stop the dump but do not release the cgroup
 		ad.Finalize(false)
 		seclog.Infof("tracing paused for [%s]", ad.GetSelectorStr())
@@ -867,9 +870,6 @@ func (adm *ActivityDumpManager) triggerLoadController() {
 			adm.emptyDropped.Inc()
 		}
 
-		// restart a new dump for the same workload
-		newDump := adm.loadController.NextPartialDump(ad)
-
 		adm.Lock()
 		if err := adm.insertActivityDump(newDump); err != nil {
 			seclog.Errorf("couldn't resume tracing [%s]: %v", newDump.GetSelectorStr(), err)
diff --git a/pkg/security/tests/fake_tags_resolver.go b/pkg/security/tests/fake_tags_resolver.go
index 45d1d9c548450..9c97502c8d266 100644
--- a/pkg/security/tests/fake_tags_resolver.go
+++ b/pkg/security/tests/fake_tags_resolver.go
@@ -40,7 +40,6 @@ func (fr *FakeResolver) Stop() error {
 func (fr *FakeResolver) Resolve(containerID string) []string {
 	fakeTags := []string{
 		"image_tag:latest",
-		"container_id:" + containerID,
 	}
 	fr.Lock()
 	defer fr.Unlock()