From f58b46cb97ff9007e4d7f5cf9994675d79060a0b Mon Sep 17 00:00:00 2001
From: David Porter <david@porter.me>
Date: Fri, 13 Jan 2023 15:31:23 -0800
Subject: [PATCH] fix process stats

Signed-off-by: David Porter <david@porter.me>
---
 pkg/kubelet/stats/cri_stats_provider.go | 17 ++++------
 pkg/kubelet/stats/helper.go             | 25 ++++++++++++++
 pkg/kubelet/stats/helper_test.go        | 43 +++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 10 deletions(-)

diff --git a/pkg/kubelet/stats/cri_stats_provider.go b/pkg/kubelet/stats/cri_stats_provider.go
index 243e8eeb590..0ecbc1d53c9 100644
--- a/pkg/kubelet/stats/cri_stats_provider.go
+++ b/pkg/kubelet/stats/cri_stats_provider.go
@@ -211,7 +211,6 @@ func (p *criStatsProvider) listPodStatsPartiallyFromCRI(ctx context.Context, upd
 		p.addPodNetworkStats(ps, podSandboxID, caInfos, cs, containerNetworkStats[podSandboxID])
 		p.addPodCPUMemoryStats(ps, types.UID(podSandbox.Metadata.Uid), allInfos, cs)
 		p.addSwapStats(ps, types.UID(podSandbox.Metadata.Uid), allInfos, cs)
-		p.addProcessStats(ps, types.UID(podSandbox.Metadata.Uid), allInfos, cs)
 
 		// If cadvisor stats is available for the container, use it to populate
 		// container stats
@@ -220,7 +219,9 @@ func (p *criStatsProvider) listPodStatsPartiallyFromCRI(ctx context.Context, upd
 			klog.V(5).InfoS("Unable to find cadvisor stats for container", "containerID", containerID)
 		} else {
 			p.addCadvisorContainerStats(cs, &caStats)
+			p.addProcessStats(ps, &caStats)
 		}
+
 		ps.Containers = append(ps.Containers, *cs)
 	}
 	// cleanup outdated caches.
@@ -584,16 +585,12 @@ func (p *criStatsProvider) addSwapStats(
 
 func (p *criStatsProvider) addProcessStats(
 	ps *statsapi.PodStats,
-	podUID types.UID,
-	allInfos map[string]cadvisorapiv2.ContainerInfo,
-	cs *statsapi.ContainerStats,
+	container *cadvisorapiv2.ContainerInfo,
 ) {
-	// try get process stats from cadvisor only.
-	info := getCadvisorPodInfoFromPodUID(podUID, allInfos)
-	if info != nil {
-		ps.ProcessStats = cadvisorInfoToProcessStats(info)
-		return
-	}
+	processStats := cadvisorInfoToProcessStats(container)
+	// Sum up all of the process stats for each of the containers to obtain the cumulative pod level process count
+	ps.ProcessStats = mergeProcessStats(ps.ProcessStats, processStats)
+	return
 }
 
 func (p *criStatsProvider) makeContainerStats(
diff --git a/pkg/kubelet/stats/helper.go b/pkg/kubelet/stats/helper.go
index c6ca3a064e0..27882df9359 100644
--- a/pkg/kubelet/stats/helper.go
+++ b/pkg/kubelet/stats/helper.go
@@ -168,6 +168,31 @@ func cadvisorInfoToProcessStats(info *cadvisorapiv2.ContainerInfo) *statsapi.Pro
 	return &statsapi.ProcessStats{ProcessCount: uint64Ptr(num)}
 }
 
+func mergeProcessStats(first *statsapi.ProcessStats, second *statsapi.ProcessStats) *statsapi.ProcessStats {
+	if first == nil && second == nil {
+		return nil
+	}
+
+	if first == nil {
+		return second
+	}
+	if second == nil {
+		return first
+	}
+
+	firstProcessCount := uint64(0)
+	if first.ProcessCount != nil {
+		firstProcessCount = *first.ProcessCount
+	}
+
+	secondProcessCount := uint64(0)
+	if second.ProcessCount != nil {
+		secondProcessCount = *second.ProcessCount
+	}
+
+	return &statsapi.ProcessStats{ProcessCount: uint64Ptr(firstProcessCount + secondProcessCount)}
+}
+
 // cadvisorInfoToNetworkStats returns the statsapi.NetworkStats converted from
 // the container info from cadvisor.
 func cadvisorInfoToNetworkStats(info *cadvisorapiv2.ContainerInfo) *statsapi.NetworkStats {
diff --git a/pkg/kubelet/stats/helper_test.go b/pkg/kubelet/stats/helper_test.go
index 38486bd9f85..050a8102c9c 100644
--- a/pkg/kubelet/stats/helper_test.go
+++ b/pkg/kubelet/stats/helper_test.go
@@ -22,10 +22,12 @@ import (
 
 	cadvisorapiv1 "github.com/google/cadvisor/info/v1"
 	cadvisorapiv2 "github.com/google/cadvisor/info/v2"
+	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/assert"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	statsapi "k8s.io/kubelet/pkg/apis/stats/v1alpha1"
+	"k8s.io/utils/pointer"
 )
 
 func TestCustomMetrics(t *testing.T) {
@@ -97,3 +99,44 @@ func TestCustomMetrics(t *testing.T) {
 			Value: 2.1,
 		})
 }
+
+func TestMergeProcessStats(t *testing.T) {
+	for _, tc := range []struct {
+		desc     string
+		first    *statsapi.ProcessStats
+		second   *statsapi.ProcessStats
+		expected *statsapi.ProcessStats
+	}{
+		{
+			desc:     "both nil",
+			first:    nil,
+			second:   nil,
+			expected: nil,
+		},
+		{
+			desc:     "first non-nil, second not",
+			first:    &statsapi.ProcessStats{ProcessCount: pointer.Uint64(100)},
+			second:   nil,
+			expected: &statsapi.ProcessStats{ProcessCount: pointer.Uint64(100)},
+		},
+		{
+			desc:     "first nil, second non-nil",
+			first:    nil,
+			second:   &statsapi.ProcessStats{ProcessCount: pointer.Uint64(100)},
+			expected: &statsapi.ProcessStats{ProcessCount: pointer.Uint64(100)},
+		},
+		{
+			desc:     "both non nill",
+			first:    &statsapi.ProcessStats{ProcessCount: pointer.Uint64(100)},
+			second:   &statsapi.ProcessStats{ProcessCount: pointer.Uint64(100)},
+			expected: &statsapi.ProcessStats{ProcessCount: pointer.Uint64(200)},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			got := mergeProcessStats(tc.first, tc.second)
+			if diff := cmp.Diff(tc.expected, got); diff != "" {
+				t.Fatalf("Unexpected diff on process stats (-want,+got):\n%s", diff)
+			}
+		})
+	}
+}