fix: prevent memory leak by reusing gRPC conn to kubelet pod-resources API

- Add persistent grpcConn to PodMapper, reused across scrapes instead of
  creating a new *grpc.ClientConn on every /metrics request (issue #702)
- Reset cached conn on RPC error to handle kubelet restarts
- Move resolver.SetDefaultScheme to one-time init in NewPodMapper
- Clone labels map per entity in expCollector and gpuHealthStatusCollector
  to prevent cross-GPU label pollution and extra GC pressure
- Add tests: connection reuse, reconnect after shutdown, Stop() cleanup

Author: AshwinRamaniPsg

internal/pkg/collector/expcollector.go

@@ -102,12 +102,9 @@ func (c *expCollector) getMetrics() (MetricsByCounter, error) {
 		entityValues, exists := mapEntityIDToValues[mi.Entity]
 		if exists {
 			for entityValue, val := range entityValues {
-
 				metricValueLabels := maps.Clone(labels)
 				c.labelFiller(metricValueLabels, entityValue)
-
 				m := c.createMetric(metricValueLabels, mi, uuid, val)
-
 				metrics[c.counter] = append(metrics[c.counter], m)
 			}
 		} else {

internal/pkg/transformation/kubernetes.go

@@ -35,6 +35,7 @@ import (
 	"google.golang.org/grpc/status"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials/insecure"
 
 	"github.com/NVIDIA/go-dcgm/pkg/dcgm"
@@ -354,6 +355,37 @@ func (p *PodMapper) Run() {
 
 func (p *PodMapper) Stop() {
 	close(p.stopChan)
+	p.grpcConnMu.Lock()
+	defer p.grpcConnMu.Unlock()
+	if p.grpcConn != nil {
+		p.grpcConn.Close()
+		p.grpcConn = nil
+	}
+}
+
+// getGRPCConn returns the cached gRPC connection to the kubelet pod-resources socket,
+// creating a new one if the cached connection is absent or has shut down.
+// Reusing the connection across scrapes is the primary fix for the RSS growth in issue #702:
+// each grpc.NewClient call allocates HTTP/2 frame buffers and spawns goroutines whose
+// cleanup is asynchronous, so creating one per scrape causes steady heap growth.
+func (p *PodMapper) getGRPCConn(socketPath string) (*grpc.ClientConn, error) {
+	p.grpcConnMu.Lock()
+	defer p.grpcConnMu.Unlock()
+
+	if p.grpcConn != nil {
+		if state := p.grpcConn.GetState(); state != connectivity.Shutdown {
+			return p.grpcConn, nil
+		}
+		p.grpcConn.Close()
+		p.grpcConn = nil
+	}
+
+	conn, _, err := connectToServer(socketPath)
+	if err != nil {
+		return nil, err
+	}
+	p.grpcConn = conn
+	return conn, nil
 }
 
 func (p *PodMapper) getMappings(deviceInfo deviceinfo.Provider) (map[string][]PodInfo, map[string]PodInfo, map[string][]PodInfo, error) {
@@ -365,14 +397,21 @@ func (p *PodMapper) getMappings(deviceInfo deviceinfo.Provider) (map[string][]Po
 		return nil, nil, nil, err
 	}
 
-	c, cleanup, err := connectToServer(socketPath)
+	c, err := p.getGRPCConn(socketPath)
 	if err != nil {
 		return nil, nil, nil, err
 	}
-	defer cleanup()
 
 	pods, err := p.listPods(c)
 	if err != nil {
+		// Reset the cached connection on RPC failure so the next scrape
+		// establishes a fresh one (e.g., after kubelet restart).
+		p.grpcConnMu.Lock()
+		if p.grpcConn != nil {
+			p.grpcConn.Close()
+			p.grpcConn = nil
+		}
+		p.grpcConnMu.Unlock()
 		return nil, nil, nil, err
 	}
 

internal/pkg/transformation/kubernetes_test.go

@@ -37,6 +37,7 @@ import (
 	"go.uber.org/mock/gomock"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/status"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -2678,3 +2679,106 @@ func TestProcessKeepsGenericWarningOnNonResourceExhaustedPodResourcesError(t *te
 	require.Contains(t, gotLog, "Failed to get pod mappings")
 	require.NotContains(t, gotLog, "Kubelet pod-resources response exceeded gRPC receive limit")
 }
+
+// TestGetGRPCConn_ReusesConnection verifies that getGRPCConn returns the same
+// *grpc.ClientConn on successive calls, avoiding the per-scrape allocation
+// overhead that caused the RSS growth in issue #702.
+func TestGetGRPCConn_ReusesConnection(t *testing.T) {
+	testutils.RequireLinux(t)
+
+	tmpDir, cleanupDir := testutils.CreateTmpDir(t)
+	defer cleanupDir()
+
+	socketPath := filepath.Join(tmpDir, "kubelet.sock")
+	server := grpc.NewServer()
+	podresourcesapi.RegisterPodResourcesListerServer(server,
+		testutils.NewMockPodResourcesServer(appconfig.NvidiaResourceName, []string{"gpu-0"}))
+	stopServer := testutils.StartMockServer(t, server, socketPath)
+	defer stopServer()
+
+	pm := &PodMapper{
+		Config:   &appconfig.Config{PodResourcesKubeletSocket: socketPath},
+		stopChan: make(chan struct{}),
+	}
+
+	conn1, err := pm.getGRPCConn(socketPath)
+	require.NoError(t, err)
+	require.NotNil(t, conn1)
+
+	conn2, err := pm.getGRPCConn(socketPath)
+	require.NoError(t, err)
+	require.NotNil(t, conn2)
+
+	assert.Same(t, conn1, conn2, "getGRPCConn must return the cached connection on repeated calls")
+}
+
+// TestGetGRPCConn_ReconnectsAfterShutdown verifies that getGRPCConn creates a
+// fresh connection when the cached one has been shut down (e.g., after a kubelet
+// restart that invalidates the Unix socket).
+func TestGetGRPCConn_ReconnectsAfterShutdown(t *testing.T) {
+	testutils.RequireLinux(t)
+
+	tmpDir, cleanupDir := testutils.CreateTmpDir(t)
+	defer cleanupDir()
+
+	socketPath := filepath.Join(tmpDir, "kubelet.sock")
+	server := grpc.NewServer()
+	podresourcesapi.RegisterPodResourcesListerServer(server,
+		testutils.NewMockPodResourcesServer(appconfig.NvidiaResourceName, []string{"gpu-0"}))
+	stopServer := testutils.StartMockServer(t, server, socketPath)
+	defer stopServer()
+
+	pm := &PodMapper{
+		Config:   &appconfig.Config{PodResourcesKubeletSocket: socketPath},
+		stopChan: make(chan struct{}),
+	}
+
+	conn1, err := pm.getGRPCConn(socketPath)
+	require.NoError(t, err)
+	require.NotNil(t, conn1)
+
+	// Forcibly shut down the cached connection to simulate a broken link.
+	conn1.Close()
+	require.Eventually(t, func() bool {
+		return conn1.GetState() == connectivity.Shutdown
+	}, 2*time.Second, 10*time.Millisecond, "connection should reach Shutdown state after Close()")
+
+	conn2, err := pm.getGRPCConn(socketPath)
+	require.NoError(t, err)
+	require.NotNil(t, conn2)
+
+	assert.NotSame(t, conn1, conn2, "getGRPCConn should allocate a new connection after the cached one shuts down")
+}
+
+// TestPodMapper_Stop_ClosesGRPCConn verifies that Stop() closes the persistent
+// gRPC connection and clears the cached pointer.
+func TestPodMapper_Stop_ClosesGRPCConn(t *testing.T) {
+	testutils.RequireLinux(t)
+
+	tmpDir, cleanupDir := testutils.CreateTmpDir(t)
+	defer cleanupDir()
+
+	socketPath := filepath.Join(tmpDir, "kubelet.sock")
+	server := grpc.NewServer()
+	podresourcesapi.RegisterPodResourcesListerServer(server,
+		testutils.NewMockPodResourcesServer(appconfig.NvidiaResourceName, []string{"gpu-0"}))
+	stopServer := testutils.StartMockServer(t, server, socketPath)
+	defer stopServer()
+
+	pm := &PodMapper{
+		Config:   &appconfig.Config{PodResourcesKubeletSocket: socketPath},
+		stopChan: make(chan struct{}),
+	}
+
+	conn, err := pm.getGRPCConn(socketPath)
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+
+	pm.Stop()
+
+	assert.Eventually(t, func() bool {
+		return conn.GetState() == connectivity.Shutdown
+	}, 2*time.Second, 10*time.Millisecond, "gRPC connection should be shut down after Stop()")
+
+	assert.Nil(t, pm.grpcConn, "grpcConn field should be nil after Stop()")
+}

internal/pkg/transformation/types.go

@@ -22,6 +22,7 @@ import (
 	"regexp"
 	"sync"
 
+	"google.golang.org/grpc"
 	"k8s.io/client-go/informers"
 	"k8s.io/client-go/kubernetes"
 	corev1listers "k8s.io/client-go/listers/core/v1"
@@ -48,6 +49,11 @@ type PodMapper struct {
 	podLister            corev1listers.PodLister
 	podInformerSynced    cache.InformerSynced
 	stopChan             chan struct{}
+	// grpcConn is a persistent connection to the kubelet pod-resources endpoint,
+	// reused across scrapes to avoid the per-scrape allocation/goroutine overhead
+	// of grpc.NewClient that causes the slow RSS growth reported in issue #702.
+	grpcConn   *grpc.ClientConn
+	grpcConnMu sync.Mutex
 }
 
 // LabelFilterCache provides efficient caching for label filtering decisions