windows: trade heap for stack to build process tree for stats in line…

…ar space (#24182)

In #20619 we overhauled how we were gathering stats for Windows
processes. Unlike in Linux where we can ask for processes in a cgroup, on
Windows we have to make a single expensive syscall to get all the processes and
then build the tree ourselves. Our algorithm to do so is recursive and quadratic
in both steps and space with the number of processes on the host. For busy hosts
this hits the stack limit and panics the Nomad client.

We already build a map of parent PID to PID, so modify this to be a map of
parent PID to slice of children and then traverse that tree only from the root
we care about (the executor PID). This moves the allocations to the heap but
makes the stats gathering linear in steps and space required.

This changeset also moves as much of this code as possible into an area
 not conditionally-compiled by OS, as the tagged test file was not being run in CI.

Fixes: #23984

drivers/shared/executor/procstats/list_test.go

@@ -0,0 +1,110 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package procstats
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/mitchellh/go-ps"
+	"github.com/shoenig/test/must"
+)
+
+type mockProcess struct {
+	pid  int
+	ppid int
+}
+
+func (p *mockProcess) Pid() int {
+	return p.pid
+}
+
+func (p *mockProcess) PPid() int {
+	return p.ppid
+}
+
+func (p *mockProcess) Executable() string {
+	return ""
+}
+
+func mockProc(pid, ppid int) *mockProcess {
+	return &mockProcess{pid: pid, ppid: ppid}
+}
+
+func genMockProcs(needles, haystack int) ([]ps.Process, []ProcessID) {
+
+	procs := []ps.Process{mockProc(1, 1), mockProc(42, 1)}
+	expect := []ProcessID{42}
+
+	// TODO: make this into a tree structure, not just a linear tree
+	for i := 0; i < needles; i++ {
+		parent := 42 + i
+		pid := parent + 1
+		procs = append(procs, mockProc(pid, parent))
+		expect = append(expect, pid)
+	}
+
+	for i := 0; i < haystack; i++ {
+		parent := 200 + i
+		pid := parent + 1
+		procs = append(procs, mockProc(pid, parent))
+	}
+
+	rand.Shuffle(len(procs), func(i, j int) {
+		procs[i], procs[j] = procs[j], procs[i]
+	})
+
+	return procs, expect
+}
+
+func Test_list(t *testing.T) {
+	cases := []struct {
+		name     string
+		needles  int
+		haystack int
+		expect   int
+	}{
+		{
+			name:     "minimal",
+			needles:  2,
+			haystack: 10,
+			expect:   16,
+		},
+		{
+			name:     "small needles small haystack",
+			needles:  5,
+			haystack: 200,
+			expect:   212,
+		},
+		{
+			name:     "small needles large haystack",
+			needles:  10,
+			haystack: 1000,
+			expect:   1022,
+		},
+		{
+			name:     "moderate needles giant haystack",
+			needles:  20,
+			haystack: 2000,
+			expect:   2042,
+		},
+	}
+
+	for _, tc := range cases {
+		const executorPID = 42
+		t.Run(tc.name, func(t *testing.T) {
+
+			procs, expect := genMockProcs(tc.needles, tc.haystack)
+			lister := func() ([]ps.Process, error) {
+				return procs, nil
+			}
+
+			result, examined := list(executorPID, lister)
+			must.SliceContainsAll(t, expect, result.Slice(),
+				must.Sprintf("exp: %v; got: %v", expect, result),
+			)
+			must.Eq(t, tc.expect, examined)
+		})
+	}
+}

drivers/shared/executor/procstats/list_windows.go

@@ -10,49 +10,6 @@ import (
 	"github.com/mitchellh/go-ps"
 )
 
-func gather(procs map[int]ps.Process, family set.Collection[int], root int, candidate ps.Process) bool {
-	if candidate == nil {
-		return false
-	}
-	pid := candidate.Pid()
-	if pid == 0 || pid == 1 {
-		return false
-	}
-	if pid == root {
-		return true
-	}
-	parent := procs[candidate.PPid()]
-	result := gather(procs, family, root, parent)
-	if result {
-		family.Insert(pid)
-	}
-	return result
-}
-
-func mapping(all []ps.Process) map[int]ps.Process {
-	result := make(map[int]ps.Process)
-	for _, process := range all {
-		result[process.Pid()] = process
-	}
-	return result
-}
-
-func list(executorPID int, processes func() ([]ps.Process, error)) set.Collection[ProcessID] {
-	family := set.From([]int{executorPID})
-
-	all, err := processes()
-	if err != nil {
-		return family
-	}
-
-	m := mapping(all)
-	for _, candidate := range all {
-		gather(m, family, executorPID, candidate)
-	}
-
-	return family
-}
-
 // List will scan the process table and return a set of the process family
 // tree starting with executorPID as the root.
 //
@@ -69,5 +26,6 @@ func list(executorPID int, processes func() ([]ps.Process, error)) set.Collectio
 // happens when you use syscalls to work your way from the root down to its
 // descendants.
 func List(executorPID int) set.Collection[ProcessID] {
-	return list(executorPID, ps.Processes)
+	procs, _ := list(executorPID, ps.Processes)
+	return procs
 }

drivers/shared/executor/procstats/procstats.go

@@ -9,6 +9,7 @@ import (
 	"github.com/hashicorp/go-set/v3"
 	"github.com/hashicorp/nomad/client/lib/cpustats"
 	"github.com/hashicorp/nomad/plugins/drivers"
+	"github.com/mitchellh/go-ps"
 )
 
 var (
@@ -80,3 +81,52 @@ func Aggregate(systemStats *cpustats.Tracker, procStats ProcUsages) *drivers.Tas
 		Pids:          procStats,
 	}
 }
+
+func list(executorPID int, processes func() ([]ps.Process, error)) (set.Collection[ProcessID], int) {
+	family := set.From([]int{executorPID})
+
+	all, err := processes()
+	if err != nil {
+		return family, 0
+	}
+
+	parents, examined := mapping(all)
+	examined += gather(family, parents, executorPID)
+
+	return family, examined
+}
+
+func gather(family set.Collection[int], parents map[int]set.Collection[int], parent int) int {
+	examined := 0
+	candidates, ok := parents[parent]
+	if !ok {
+		return examined
+	}
+	for _, candidate := range candidates.Slice() {
+		examined++
+		family.Insert(candidate)
+		examined += gather(family, parents, candidate)
+	}
+
+	return examined
+}
+
+// mapping builds a reverse map of parent to children
+func mapping(all []ps.Process) (map[int]set.Collection[int], int) {
+
+	parents := map[int]set.Collection[int]{}
+	examined := 0
+
+	for _, candidate := range all {
+		if candidate != nil {
+			examined++
+			if children, ok := parents[candidate.PPid()]; ok {
+				children.Insert(candidate.Pid())
+			} else {
+				parents[candidate.PPid()] = set.From([]int{candidate.Pid()})
+			}
+		}
+	}
+
+	return parents, examined
+}