fix: return deep copies in GetScheduledPods to prevent data races

cmd/scheduler/metrics.go

@@ -20,7 +20,6 @@ import (
 	"fmt"
 	"log"
 	"net/http"
-	"strings"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -355,25 +354,6 @@ func (cc ClusterManagerCollector) Collect(ch chan<- prometheus.Metric) {
 							float64(ctrdevval.Usedcores),
 							val.Namespace, val.NodeID, val.Name, fmt.Sprint(ctridx), ctrdevval.UUID)
 					}
-					var totaldev int32
-					found := false
-					for _, ni := range *nu {
-						for _, nodedev := range ni.Devices.DeviceLists {
-							if strings.Compare(nodedev.Device.ID, ctrdevval.UUID) == 0 {
-								totaldev = nodedev.Device.Totalmem
-								found = true
-								break
-							}
-						}
-						if found {
-							break
-						}
-					}
-					klog.V(4).InfoS("Total memory for device",
-						"deviceUUID", ctrdevval.UUID,
-						"totalMemory", totaldev,
-						"nodeID", val.NodeID,
-					)
 				}
 			}
 		}

pkg/device/pod_test.go

@@ -18,7 +18,6 @@ package device
 
 import (
 	"reflect"
-	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -122,57 +121,76 @@ func TestPodUseDeviceStat(t *testing.T) {
 		})
 	}
 }
-func TestGetScheduledPods(t *testing.T) {
-	podManager := &PodManager{
-		pods:  make(map[k8stypes.UID]*PodInfo),
-		mutex: sync.RWMutex{},
-	}
+func TestGetScheduledPodsReturnsDeepCopy(t *testing.T) {
+	podManager := NewPodManager()
 
-	pod1 := &PodInfo{
-		Pod: &corev1.Pod{
-			ObjectMeta: metav1.ObjectMeta{
-				Namespace: "default",
-				Name:      "pod1",
-				UID:       k8stypes.UID("uid1"),
-			},
+	pod1 := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Namespace: "default",
+			Name:      "pod1",
+			UID:       k8stypes.UID("uid1"),
 		},
-		NodeID:  "node1",
-		Devices: PodDevices{"device1": {{}}},
 	}
-	pod2 := &PodInfo{
-		Pod: &corev1.Pod{
-			ObjectMeta: metav1.ObjectMeta{
-				Namespace: "default",
-				Name:      "pod2",
-				UID:       k8stypes.UID("uid2"),
+	pod1Devices := PodDevices{
+		"NVIDIA": {
+			{
+				{
+					Idx:       0,
+					UUID:      "GPU-1",
+					Type:      "NVIDIA",
+					Usedmem:   1000,
+					Usedcores: 50,
+					CustomInfo: map[string]any{
+						"annotations": map[string]string{
+							"metax.com/gpu": "true",
+						},
+					},
+				},
 			},
 		},
-
-		NodeID:  "node2",
-		Devices: PodDevices{"device2": {{}}},
 	}
-	podManager.pods[pod1.UID] = pod1
-	podManager.pods[pod2.UID] = pod2
+
+	podManager.AddPod(pod1, "node1", pod1Devices)
 
 	scheduledPods, err := podManager.GetScheduledPods()
 
 	assert.NoError(t, err, "GetScheduledPods should not return an error")
 	assert.NotNil(t, scheduledPods, "The result should not be nil")
-	assert.Equal(t, 2, len(scheduledPods), "The number of scheduled pods should be 2")
+	assert.Equal(t, 1, len(scheduledPods), "The number of scheduled pods should be 1")
+
+	got, ok := scheduledPods[pod1.UID]
+	assert.True(t, ok)
+
+	// 1. Existing Pod pointer is kept (retaining pointer is intentional)
+	assert.Same(t, pod1, got.Pod, "Pod pointer should be preserved without calling Pod.DeepCopy()")
+	assert.Equal(t, "node1", got.NodeID)
+
+	// 2. Scalar device allocation fields match
+	gotDev := got.Devices["NVIDIA"][0][0]
+	assert.Equal(t, "GPU-1", gotDev.UUID)
+	assert.Equal(t, "NVIDIA", gotDev.Type)
+	assert.Equal(t, int32(1000), gotDev.Usedmem)
+	assert.Equal(t, int32(50), gotDev.Usedcores)
+
+	// 3. CustomInfo is intentionally omitted (nil) in metrics snapshot
+	assert.Nil(t, gotDev.CustomInfo, "CustomInfo should be nil in metrics snapshot")
+
+	// 4. Device allocation fields are independent; mutating snapshot does not affect PodManager
+	got.Devices["NVIDIA"][0][0].UUID = "MUTATED-GPU"
+	got.Devices["NVIDIA"][0][0].Usedmem = 9999
+	got.Devices["NVIDIA"][0][0].Usedcores = 99
+
+	originalInfo, ok := podManager.GetPod(pod1)
+	assert.True(t, ok)
+	origDev := originalInfo.Devices["NVIDIA"][0][0]
+	assert.Equal(t, "GPU-1", origDev.UUID, "Original UUID should remain unmutated")
+	assert.Equal(t, int32(1000), origDev.Usedmem, "Original Usedmem should remain unmutated")
+	assert.Equal(t, int32(50), origDev.Usedcores, "Original Usedcores should remain unmutated")
+	assert.NotNil(t, origDev.CustomInfo, "Original CustomInfo should remain present in PodManager")
+}
 
-	expectedPods := map[k8stypes.UID]*PodInfo{
-		pod1.UID: pod1,
-		pod2.UID: pod2,
-	}
-	for uid, pod := range scheduledPods {
-		expectedPod := expectedPods[uid]
-		assert.NotNil(t, expectedPod, "Pod with UID %s should exist in the expected pods", uid)
-		assert.Equal(t, expectedPod.Namespace, pod.Namespace, "Namespace should match")
-		assert.Equal(t, expectedPod.Name, pod.Name, "Name should match")
-		assert.Equal(t, expectedPod.UID, pod.UID, "UID should match")
-		assert.Equal(t, expectedPod.NodeID, pod.NodeID, "NodeID should match")
-		assert.Equal(t, expectedPod.Devices, pod.Devices, "Devices should match")
-	}
+func TestGetScheduledPods(t *testing.T) {
+	TestGetScheduledPodsReturnsDeepCopy(t)
 }
 
 func TestGetPod(t *testing.T) {
@@ -516,16 +534,19 @@ func TestContainerDeviceDeepCopy(t *testing.T) {
 
 	copy := original.DeepCopy()
 
-	// 1. Copy must be deeply equal to original.
-	assert.Equal(t, original, copy)
+	// 1. Scalar fields match original.
+	assert.Equal(t, original.Idx, copy.Idx)
+	assert.Equal(t, original.UUID, copy.UUID)
+	assert.Equal(t, original.Type, copy.Type)
+	assert.Equal(t, original.Usedmem, copy.Usedmem)
+	assert.Equal(t, original.Usedcores, copy.Usedcores)
 
-	// 2. Mutating the copy must not affect the original.
-	copy.UUID = "mutated-gpu"
-	copy.CustomInfo["key2"] = "value2"
+	// 2. CustomInfo is intentionally omitted (nil).
+	assert.Nil(t, copy.CustomInfo, "CustomInfo should be intentionally omitted in DeepCopy")
 
-	assert.Equal(t, original.UUID, "GPU-0")
-	_, exists := original.CustomInfo["key2"]
-	assert.False(t, exists, "original CustomInfo should not have key2")
+	// 3. Mutating scalar fields of the copy does not affect the original.
+	copy.UUID = "mutated-gpu"
+	assert.Equal(t, "GPU-0", original.UUID)
 }
 
 func TestListPodsInfoReturnsDeepCopy(t *testing.T) {

pkg/device/pods.go

@@ -17,7 +17,6 @@ limitations under the License.
 package device
 
 import (
-	"maps"
 	"sync"
 
 	corev1 "k8s.io/api/core/v1"
@@ -212,16 +211,32 @@ func (cd ContainerDevices) DeepCopy() ContainerDevices {
 }
 
 func (c ContainerDevice) DeepCopy() ContainerDevice {
-	dup := ContainerDevice{
+	return ContainerDevice{
 		Idx:       c.Idx,
 		UUID:      c.UUID,
 		Type:      c.Type,
 		Usedmem:   c.Usedmem,
 		Usedcores: c.Usedcores,
 	}
-	if c.CustomInfo != nil {
-		dup.CustomInfo = make(map[string]any, len(c.CustomInfo))
-		maps.Copy(dup.CustomInfo, c.CustomInfo)
+}
+
+func (pd PodDevices) DeepCopyForMetrics() PodDevices {
+	if pd == nil {
+		return nil
+	}
+
+	dup := make(PodDevices, len(pd))
+	for deviceType, podSingleDevice := range pd {
+		deviceCopy := make(PodSingleDevice, len(podSingleDevice))
+		for containerIndex, containerDevices := range podSingleDevice {
+			deviceCopy[containerIndex] = make(ContainerDevices, len(containerDevices))
+			copy(deviceCopy[containerIndex], containerDevices)
+
+			for i := range deviceCopy[containerIndex] {
+				deviceCopy[containerIndex][i].CustomInfo = nil
+			}
+		}
+		dup[deviceType] = deviceCopy
 	}
 	return dup
 }
@@ -230,14 +245,13 @@ func (m *PodManager) GetScheduledPods() (map[k8stypes.UID]*PodInfo, error) {
 	m.mutex.RLock()
 	defer m.mutex.RUnlock()
 
-	podCount := len(m.pods)
-	klog.InfoS("Retrieved scheduled pods",
-		"podCount", podCount,
-	)
-
-	// Return a shallow copy of the pods map to avoid race conditions.
-	// This prevents a "concurrent map iteration and map write" fatal error.
-	podsCopy := make(map[k8stypes.UID]*PodInfo, podCount)
-	maps.Copy(podsCopy, m.pods)
+	podsCopy := make(map[k8stypes.UID]*PodInfo, len(m.pods))
+	for uid, pod := range m.pods {
+		podsCopy[uid] = &PodInfo{
+			Pod:     pod.Pod,
+			NodeID:  pod.NodeID,
+			Devices: pod.Devices.DeepCopyForMetrics(),
+		}
+	}
 	return podsCopy, nil
 }