fix(processtree): walk container branches under the creator lock

pkg/processtree/creator/processtree_creator.go

@@ -148,7 +148,8 @@ func (pt *processTreeCreatorImpl) GetProcessBootTimeNs(pid uint32) uint64 {
 	return pt.pidStartTimeNs[pid]
 }
 
-// GetPidBranch performs container branch operation (no longer needs to be atomic)
+// GetPidBranch walks the container branch under the creator's lock, so exit
+// cleanup cannot reparent nodes mid-walk.
 func (pt *processTreeCreatorImpl) GetPidBranch(containerTree any, containerID string, targetPID uint32) (armotypes.Process, error) {
 	pt.mutex.RLock()
 	defer pt.mutex.RUnlock()

pkg/processtree/creator/processtree_creator_interface.go

@@ -19,6 +19,8 @@ type ProcessTreeCreator interface {
 	// since boot (CLOCK_BOOTTIME), sourced exclusively from /proc/<pid>/stat
 	// field 22. Returns 0 when unknown (process not yet seen, or already exited).
 	GetProcessBootTimeNs(pid uint32) uint64
+	// GetPidBranch returns the container branch for a pid under the creator's lock
+	GetPidBranch(containerTree any, containerID string, targetPID uint32) (armotypes.Process, error)
 	// Start the process tree creator and begin background tasks
 	Start()
 	// Stop the process tree creator and cleanup resources

pkg/processtree/process_tree_manager.go

@@ -87,7 +87,7 @@ func (ptm *ProcessTreeManagerImpl) GetContainerProcessTree(containerID string, p
 	func() {
 		ptm.mutex.RLock()
 		defer ptm.mutex.RUnlock()
-		containerSubtree, subtreeErr = ptm.containerTree.GetPidBranch(containerID, pid, ptm.creator.GetProcessMap())
+		containerSubtree, subtreeErr = ptm.creator.GetPidBranch(ptm.containerTree, containerID, pid)
 	}()
 
 	if subtreeErr != nil {

pkg/processtree/process_tree_manager_test.go

@@ -3,6 +3,7 @@ package processtree
 import (
 	"fmt"
 	"os"
+	"sync"
 	"testing"
 	"time"
 
@@ -144,3 +145,46 @@ func TestManager_GetContainerProcessTree_CarriesStartTime(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, "nginx", bypassed.Comm)
 }
+
+func TestManager_GetContainerProcessTree_ConcurrentExitCleanup(t *testing.T) {
+	cfg := config.Config{ExitCleanup: treeconfig.ExitCleanupConfig{
+		MaxPendingExits: 100, CleanupInterval: time.Millisecond, CleanupDelay: 0,
+	}}
+	containerTree := containerprocesstree.NewContainerProcessTree()
+	creator := processtreecreator.NewProcessTreeCreator(containerTree, cfg)
+	mgr := NewProcessTreeManager(creator, containerTree, cfg)
+	containerTree.ContainerCallback(containercollection.PubSubEvent{
+		Type: containercollection.EventTypeAddContainer,
+		Container: &containercollection.Container{
+			Runtime: containercollection.RuntimeMetadata{
+				BasicRuntimeMetadata: eventtypes.BasicRuntimeMetadata{ContainerID: "host", ContainerPID: 1},
+			},
+		},
+	})
+	require.NoError(t, mgr.ReportEvent(utils.ProcfsEventType, &events.ProcfsEvent{PID: 1, Comm: "systemd", ContainerID: "host"}))
+	mgr.Start()
+	defer mgr.Stop()
+
+	const parent, child uint32 = 1 << 30, 1<<30 + 1
+	stop := make(chan struct{})
+	var wg sync.WaitGroup
+	wg.Go(func() {
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				_, _ = mgr.GetContainerProcessTree("host", child, false)
+			}
+		}
+	})
+	for range 200 {
+		now := time.Now().UnixNano()
+		require.NoError(t, mgr.ReportEvent(utils.ExecveEventType, &utils.StructEvent{ContainerID: "host", Pid: parent, Ppid: 1, Comm: "sh", Timestamp: now}))
+		require.NoError(t, mgr.ReportEvent(utils.ExecveEventType, &utils.StructEvent{ContainerID: "host", Pid: child, Ppid: parent, Comm: "sleep", Timestamp: now}))
+		require.NoError(t, mgr.ReportEvent(utils.ExitEventType, &utils.StructEvent{ContainerID: "host", Pid: parent, Ppid: 1, Timestamp: now}))
+		time.Sleep(2 * time.Millisecond)
+	}
+	close(stop)
+	wg.Wait()
+}