fix(nodelock): let a pod re-acquire its own still-valid node lock

pkg/util/nodelock/nodelock.go

@@ -233,9 +233,18 @@ func LockNode(nodeName string, lockname string, pods *corev1.Pod) error {
 	if time.Since(lockTime) > NodeLockTimeout {
 		klog.InfoS("Node lock expired", "node", nodeName, "lockTime", lockTime, "timeout", NodeLockTimeout)
 		skipOwnerCheck = true
+	} else if ns == pods.Namespace && previousPodName == pods.Name {
+		// The lock is already held by this exact pod. lockAllDevices calls
+		// LockNode once per device vendor a pod requests resources from, so
+		// a pod requesting resources from two or more vendors (e.g. both
+		// nvidia.com/gpu and cambricon.com/vmlu) would otherwise contend
+		// with its own still-valid lock on the second call and never
+		// become schedulable. Treat this as already acquired.
+		klog.V(4).InfoS("Node lock already held by this pod, treating as acquired", "node", nodeName, "podName", pods.Name)
+		return nil
 	} else
 	// Check dangling nodeLock
-	if ns != "" && previousPodName != "" && (ns != pods.Namespace || previousPodName != pods.Name) {
+	if ns != "" && previousPodName != "" {
 		if _, err := client.GetClient().CoreV1().Pods(ns).Get(ctx, previousPodName, metav1.GetOptions{}); err != nil {
 			if !apierrors.IsNotFound(err) {
 				klog.ErrorS(err, "Failed to get pod of NodeLock", "podName", previousPodName, "namespace", ns)

pkg/util/nodelock/nodelock_test.go

@@ -34,7 +34,7 @@ import (
 func Test_LockNode(t *testing.T) {
 	client.KubeClient = fake.NewClientset()
 	type args struct {
-		nodeName func() string
+		nodeName func(t *testing.T) string
 		lockname string
 		pods     *corev1.Pod
 	}
@@ -46,7 +46,7 @@ func Test_LockNode(t *testing.T) {
 		{
 			name: "node not found",
 			args: args{
-				nodeName: func() string {
+				nodeName: func(t *testing.T) string {
 					return "node"
 				},
 				pods: &corev1.Pod{
@@ -59,20 +59,67 @@ func Test_LockNode(t *testing.T) {
 			wantErr: true,
 		},
 		{
-			name: "node has been locked",
+			name: "node has been locked by another pod",
 			args: args{
-				nodeName: func() string {
+				nodeName: func(t *testing.T) string {
 					name := "worker-1"
-					client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
+					if _, err := client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
 						ObjectMeta: metav1.ObjectMeta{
 							Name: name,
 							Annotations: map[string]string{
 								NodeLockKey: GenerateNodeLockKeyByPod(&corev1.Pod{
-									ObjectMeta: metav1.ObjectMeta{Name: "hami", Namespace: "hami-ns"},
+									ObjectMeta: metav1.ObjectMeta{Name: "other-pod", Namespace: "other-ns"},
 								}),
 							},
 						},
-					}, metav1.CreateOptions{})
+					}, metav1.CreateOptions{}); err != nil {
+						t.Fatalf("failed to create node fixture: %v", err)
+					}
+					// The lock holder ("other-pod"/"other-ns") must exist and not be
+					// dangling, otherwise LockNode treats it as stale and takes over.
+					if _, err := client.KubeClient.CoreV1().Pods("other-ns").Create(context.TODO(), &corev1.Pod{
+						ObjectMeta: metav1.ObjectMeta{Name: "other-pod", Namespace: "other-ns"},
+					}, metav1.CreateOptions{}); err != nil {
+						t.Fatalf("failed to create lock-holder pod fixture: %v", err)
+					}
+					return name
+				},
+				pods: &corev1.Pod{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "hami",
+						Namespace: "hami-ns",
+					},
+				},
+			},
+			wantErr: true,
+		},
+		{
+			name: "node has been locked by another pod in the same namespace",
+			args: args{
+				nodeName: func(t *testing.T) string {
+					name := "worker-1b"
+					if _, err := client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
+						ObjectMeta: metav1.ObjectMeta{
+							Name: name,
+							Annotations: map[string]string{
+								NodeLockKey: GenerateNodeLockKeyByPod(&corev1.Pod{
+									ObjectMeta: metav1.ObjectMeta{Name: "other-pod-same-ns", Namespace: "hami-ns"},
+								}),
+							},
+						},
+					}, metav1.CreateOptions{}); err != nil {
+						t.Fatalf("failed to create node fixture: %v", err)
+					}
+					// Same namespace as the requester below, but a different pod
+					// name: exercises ns == pods.Namespace (true) with
+					// previousPodName == pods.Name (false), distinct from both the
+					// "another pod" case above (both false) and the reentrant
+					// same-pod case (both true).
+					if _, err := client.KubeClient.CoreV1().Pods("hami-ns").Create(context.TODO(), &corev1.Pod{
+						ObjectMeta: metav1.ObjectMeta{Name: "other-pod-same-ns", Namespace: "hami-ns"},
+					}, metav1.CreateOptions{}); err != nil {
+						t.Fatalf("failed to create lock-holder pod fixture: %v", err)
+					}
 					return name
 				},
 				pods: &corev1.Pod{
@@ -87,16 +134,18 @@ func Test_LockNode(t *testing.T) {
 		{
 			name: "node lock is invalid",
 			args: args{
-				nodeName: func() string {
+				nodeName: func(t *testing.T) string {
 					name := "worker-2"
-					client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
+					if _, err := client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
 						ObjectMeta: metav1.ObjectMeta{
 							Name: name,
 							Annotations: map[string]string{
 								NodeLockKey: "lock",
 							},
 						},
-					}, metav1.CreateOptions{})
+					}, metav1.CreateOptions{}); err != nil {
+						t.Fatalf("failed to create node fixture: %v", err)
+					}
 					return name
 				},
 				pods: &corev1.Pod{
@@ -111,11 +160,13 @@ func Test_LockNode(t *testing.T) {
 		{
 			name: "successfully set node lock",
 			args: args{
-				nodeName: func() string {
+				nodeName: func(t *testing.T) string {
 					name := "worker-3"
-					client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
+					if _, err := client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
 						ObjectMeta: metav1.ObjectMeta{Name: name, Annotations: map[string]string{}},
-					}, metav1.CreateOptions{})
+					}, metav1.CreateOptions{}); err != nil {
+						t.Fatalf("failed to create node fixture: %v", err)
+					}
 					return name
 				},
 				pods: &corev1.Pod{
@@ -130,7 +181,7 @@ func Test_LockNode(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if err := LockNode(tt.args.nodeName(), tt.args.lockname, tt.args.pods); (err != nil) != tt.wantErr {
+			if err := LockNode(tt.args.nodeName(t), tt.args.lockname, tt.args.pods); (err != nil) != tt.wantErr {
 				t.Errorf("LockNode() error = %v, wantErr %v", err, tt.wantErr)
 			}
 		})
@@ -233,6 +284,40 @@ func TestLockNodeWithDangling(t *testing.T) {
 	}
 }
 
+// TestLockNodeReentrantSamePod covers lockAllDevices' actual call pattern:
+// it calls LockNode once per device vendor a pod requests resources from, so
+// a pod requesting e.g. both nvidia.com/gpu and cambricon.com/vmlu locks the
+// same node twice for itself in a row. The second call must succeed instead
+// of contending with the pod's own still-valid lock, or such a pod could
+// never become schedulable.
+func TestLockNodeReentrantSamePod(t *testing.T) {
+	client.KubeClient = fake.NewClientset()
+	nodeLocks = newNodeLockManager()
+	nodeName := "multi-vendor-node"
+	_, err := client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{Name: nodeName, Annotations: map[string]string{}},
+	}, metav1.CreateOptions{})
+	if err != nil {
+		t.Fatalf("Failed to create node: %v", err)
+	}
+	pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "multi-vendor-pod", Namespace: "test-ns"}}
+	// The requesting pod must actually exist for this test to exercise the
+	// intended live-ownership scenario. Without it, the first LockNode call
+	// still succeeds, but only because it's setting a fresh lock, not
+	// because the reentrancy branch has verified anything about a real,
+	// live pod - so a real dangling-lock path could be masked instead.
+	if _, err := client.KubeClient.CoreV1().Pods("test-ns").Create(context.TODO(), pod, metav1.CreateOptions{}); err != nil {
+		t.Fatalf("Failed to create pod: %v", err)
+	}
+
+	if err := LockNode(nodeName, "nvidia", pod); err != nil {
+		t.Fatalf("first LockNode call (simulating the nvidia backend) failed: %v", err)
+	}
+	if err := LockNode(nodeName, "cambricon", pod); err != nil {
+		t.Fatalf("second LockNode call for the same pod (simulating the cambricon backend) should succeed, got: %v", err)
+	}
+}
+
 func TestReleaseNodeLock(t *testing.T) {
 	client.KubeClient = fake.NewClientset()
 	type args struct {