fix: add nodeDiskAvailabilityMap

pkg/controller/attach_detach.go

@@ -113,6 +113,8 @@ func (r *ReconcileAttachDetach) Reconcile(ctx context.Context, request reconcile
 		if err := r.removeFinalizer(ctx, azVolumeAttachment); err != nil {
 			return reconcileReturnOnError(ctx, azVolumeAttachment, "delete", err, r.retryInfo)
 		}
+		// deletion of azVolumeAttachment is succeeded, the node's remaining capacity of disk attachment should be increased by 1
+		r.incrementAttachmentCount(ctx, azVolumeAttachment.Spec.NodeName)
 		// detachment request
 	} else if volumeDetachRequested(azVolumeAttachment) {
 		if err := r.triggerDetach(ctx, azVolumeAttachment); err != nil {
@@ -282,6 +284,8 @@ func (r *ReconcileAttachDetach) triggerAttach(ctx context.Context, azVolumeAttac
 						r.eventRecorder.Eventf(pod.DeepCopyObject(), v1.EventTypeNormal, consts.ReplicaAttachmentSuccessEvent, "Replica mount for volume %s successfully attached to node %s", azVolumeAttachment.Spec.VolumeName, azVolumeAttachment.Spec.NodeName)
 					}
 				}
+				// the node's remaining capacity of disk attachment should be decreased by 1, since the disk attachment is succeeded.
+				r.decrementAttachmentCount(ctx, azVolumeAttachment.Spec.NodeName)
 			}
 
 			updateFunc := func(obj client.Object) error {

pkg/controller/attach_detach_test.go

@@ -98,6 +98,7 @@ func TestAttachDetachReconcile(t *testing.T) {
 					newAttachment)
 
 				controller.azVolumeAttachmentToVaMap.Store(newAttachment.Name, newVolumeAttachment.Name)
+				addTestNodeInAvailableAttachmentsMap(*controller.SharedState, newAttachment.Spec.NodeName)
 
 				mockClientsAndAttachmentProvisioner(controller)
 
@@ -137,6 +138,7 @@ func TestAttachDetachReconcile(t *testing.T) {
 					newVolumeAttachment,
 					newAttachment)
 
+				addTestNodeInAvailableAttachmentsMap(*controller.SharedState, newAttachment.Spec.NodeName)
 				mockClientsAndAttachmentProvisioner(controller)
 
 				return controller
@@ -168,6 +170,7 @@ func TestAttachDetachReconcile(t *testing.T) {
 					&testAzVolume0,
 					newAttachment)
 
+				addTestNodeInAvailableAttachmentsMap(*controller.SharedState, newAttachment.Spec.NodeName)
 				mockClientsAndAttachmentProvisioner(controller)
 
 				return controller

pkg/controller/azdrivernode.go

@@ -55,8 +55,12 @@ func (r *ReconcileAzDriverNode) Reconcile(ctx context.Context, request reconcile
 	n := &corev1.Node{}
 	err := r.cachedClient.Get(ctx, request.NamespacedName, n)
 
-	// If the node still exists don't delete the AzDriverNode
 	if err == nil {
+		if n.ObjectMeta.DeletionTimestamp == nil {
+			// for create event, add the new node in availableAttachmentsMap
+			r.addNodeToAvailableAttachmentsMap(ctx, request.Name)
+		}
+		// for delete even, if the node still exists don't delete the AzDriverNode
 		return reconcile.Result{}, nil
 	}
 
@@ -70,6 +74,7 @@ func (r *ReconcileAzDriverNode) Reconcile(ctx context.Context, request reconcile
 		if err != nil && !errors.IsNotFound(err) {
 			return reconcile.Result{Requeue: true}, err
 		}
+		r.deleteNodeFromAvailableAttachmentsMap(ctx, request.Name)
 
 		// Delete all volumeAttachments attached to this node, if failed, requeue
 		if _, err = r.cleanUpAzVolumeAttachmentByNode(ctx, request.Name, azdrivernode, azureutils.AllRoles, cleanUpAttachment); err != nil {
@@ -101,6 +106,7 @@ func (r *ReconcileAzDriverNode) Recover(ctx context.Context) error {
 		if _, err = r.azClient.DiskV1beta2().AzDriverNodes(r.objectNamespace).Update(ctx, updated, metav1.UpdateOptions{}); err != nil {
 			return err
 		}
+		r.addNodeToAvailableAttachmentsMap(ctx, node.Name)
 	}
 	return nil
 }
@@ -127,7 +133,7 @@ func NewAzDriverNodeController(mgr manager.Manager, controllerSharedState *Share
 	// Predicate to only reconcile deleted nodes
 	p := predicate.Funcs{
 		CreateFunc: func(e event.CreateEvent) bool {
-			return false
+			return true
 		},
 		UpdateFunc: func(e event.UpdateEvent) bool {
 			return false

pkg/controller/azdrivernode_test.go

@@ -26,6 +26,7 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/selection"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/wait"
 	"k8s.io/klog/v2/klogr"
 	azdiskfakes "sigs.k8s.io/azuredisk-csi-driver/pkg/apis/client/clientset/versioned/fake"
@@ -131,9 +132,12 @@ func TestAzDriverNodeControllerReconcile(t *testing.T) {
 					&testAzDriverNode1)
 
 				controller.cachedClient.(*mockclient.MockClient).EXPECT().
-					Get(gomock.Any(), testNode1Request.NamespacedName, gomock.Any()).
-					Return(nil).
-					AnyTimes()
+					Get(gomock.Any(), gomock.Any(), gomock.Any()).
+					Return(nil).AnyTimes()
+
+				controller.cachedClient.(*mockclient.MockClient).EXPECT().
+					List(gomock.Any(), gomock.Any(), gomock.Any()).
+					Return(nil).AnyTimes()
 
 				return controller
 			},
@@ -170,6 +174,34 @@ func TestAzDriverNodeControllerReconcile(t *testing.T) {
 				require.NoError(t, err2)
 			},
 		},
+		{
+			description: "[Success] Should add new AzDriverNode in availableAttachmentsMap",
+			request:     testNode1Request,
+			setupFunc: func(t *testing.T, mockCtl *gomock.Controller) *ReconcileAzDriverNode {
+				controller := NewTestAzDriverNodeController(
+					mockCtl,
+					testNamespace,
+					&testAzDriverNode0)
+
+				controller.cachedClient.(*mockclient.MockClient).EXPECT().
+					Get(gomock.Any(), gomock.Any(), gomock.Any()).
+					Return(nil).
+					AnyTimes()
+
+				mockClients(controller.cachedClient.(*mockclient.MockClient), controller.azClient, nil)
+
+				return controller
+			},
+			verifyFunc: func(t *testing.T, controller *ReconcileAzDriverNode, result reconcile.Result, err error) {
+				require.NoError(t, err)
+				require.False(t, result.Requeue)
+
+				_, err2 := controller.azClient.DiskV1beta2().AzDriverNodes(testNamespace).Get(context.TODO(), testNode0Name, metav1.GetOptions{})
+				require.NoError(t, err2)
+				_, nodeExists := controller.availableAttachmentsMap.Load(testNode1Name)
+				require.True(t, nodeExists)
+			},
+		},
 	}
 
 	for _, test := range tests {
@@ -200,8 +232,18 @@ func TestAzDriverNodeRecover(t *testing.T) {
 					&testAzDriverNode1)
 
 				controller.cachedClient.(*mockclient.MockClient).EXPECT().
-					Get(gomock.Any(), testNode1Request.NamespacedName, gomock.Any()).
-					Return(testNode1ServerTimeoutError).
+					Get(gomock.Any(), types.NamespacedName{Name: testAzDriverNode0.Name}, gomock.Any()).
+					Return(nil).
+					AnyTimes()
+
+				controller.cachedClient.(*mockclient.MockClient).EXPECT().
+					Get(gomock.Any(), types.NamespacedName{Name: testAzDriverNode1.Name}, gomock.Any()).
+					Return(nil).
+					AnyTimes()
+
+				controller.cachedClient.(*mockclient.MockClient).EXPECT().
+					List(gomock.Any(), gomock.Any(), gomock.Any()).
+					Return(nil).
 					AnyTimes()
 
 				return controller
@@ -225,6 +267,10 @@ func TestAzDriverNodeRecover(t *testing.T) {
 			defer mockCtl.Finish()
 			controller := tt.setupFunc(t, mockCtl)
 			err := controller.Recover(context.TODO())
+			_, node0Exists := controller.availableAttachmentsMap.Load(testAzDriverNode0.Name)
+			require.True(t, node0Exists)
+			_, node1Exists := controller.availableAttachmentsMap.Load(testAzDriverNode1.Name)
+			require.True(t, node1Exists)
 			tt.verifyFunc(t, controller, err)
 		})
 	}

pkg/controller/common_test.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/golang/mock/gomock"
@@ -751,3 +752,9 @@ func mockClients(mockClient *mockclient.MockClient, azVolumeClient azdisk.Interf
 		}).
 		AnyTimes()
 }
+
+func addTestNodeInAvailableAttachmentsMap(sharedState SharedState, nodeName string) {
+	var count atomic.Int32
+	count.Store(int32(8))
+	sharedState.availableAttachmentsMap.Store(nodeName, &count)
+}

pkg/controller/node_availability_test.go

@@ -18,6 +18,7 @@ package controller
 
 import (
 	"context"
+	"sync/atomic"
 	"testing"
 
 	"github.com/golang/mock/gomock"
@@ -81,6 +82,10 @@ func TestNodeAvailabilityController(t *testing.T) {
 					newPod,
 				)
 
+				var count atomic.Int32
+				count.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testSchedulableNode1.Name, &count)
+
 				mockClients(controller.cachedClient.(*mockclient.MockClient), controller.azClient, controller.kubeClient)
 				controller.priorityReplicaRequestsQueue.Push(context.TODO(), &ReplicaRequest{VolumeName: testPersistentVolume0Name, Priority: 1})
 

pkg/controller/pod_test.go

@@ -18,6 +18,7 @@ package controller
 
 import (
 	"context"
+	"sync/atomic"
 	"testing"
 
 	"github.com/Azure/go-autorest/autorest/to"
@@ -79,6 +80,8 @@ func TestPodReconcile(t *testing.T) {
 					&testNode0,
 				)
 
+				addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode0.Name)
+
 				mockClients(controller.cachedClient.(*mockclient.MockClient), controller.azClient, controller.kubeClient)
 				return controller
 			},
@@ -123,6 +126,14 @@ func TestPodReconcile(t *testing.T) {
 					&testPersistentVolume0,
 					newPod)
 
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode0.Name)
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode1.Name)
+				var count atomic.Int32
+				count.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode0.Name, &count)
+				var count1 atomic.Int32
+				count1.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode1.Name, &count1)
 				mockClients(controller.cachedClient.(*mockclient.MockClient), controller.azClient, controller.kubeClient)
 				return controller
 			},
@@ -177,6 +188,18 @@ func TestPodReconcile(t *testing.T) {
 					&testNode1,
 					&testNode2,
 					newPod)
+				var count atomic.Int32
+				count.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode0.Name, &count)
+				var count1 atomic.Int32
+				count1.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode1.Name, &count1)
+				var count2 atomic.Int32
+				count2.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode2.Name, &count2)
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode0.Name)
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode1.Name)
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode2.Name)
 
 				mockClients(controller.cachedClient.(*mockclient.MockClient), controller.azClient, controller.kubeClient)
 				return controller
@@ -240,6 +263,18 @@ func TestPodReconcile(t *testing.T) {
 					newPod0,
 					newPod1)
 
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode0.Name)
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode1.Name)
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode2.Name)
+				var count atomic.Int32
+				count.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode0.Name, &count)
+				var count1 atomic.Int32
+				count1.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode1.Name, &count1)
+				var count2 atomic.Int32
+				count2.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode2.Name, &count2)
 				mockClients(controller.cachedClient.(*mockclient.MockClient), controller.azClient, controller.kubeClient)
 				result, err := controller.Reconcile(context.TODO(), testPod0Request)
 				require.False(t, result.Requeue)
@@ -289,6 +324,9 @@ func TestPodReconcile(t *testing.T) {
 					&testPersistentVolume0,
 					newPod0)
 
+				addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode0.Name)
+				addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode1.Name)
+
 				mockClients(controller.cachedClient.(*mockclient.MockClient), controller.azClient, controller.kubeClient)
 
 				result, err := controller.Reconcile(context.TODO(), testPod0Request)
@@ -375,6 +413,14 @@ func TestPodRecover(t *testing.T) {
 					newAttachment1,
 					newPod)
 
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode0.Name)
+				// addTestNodeInAvailableAttachmentsMap(*controller.SharedState, testNode1.Name)
+				var count atomic.Int32
+				count.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode0.Name, &count)
+				var count1 atomic.Int32
+				count1.Store(int32(8))
+				controller.availableAttachmentsMap.Store(testNode1.Name, &count1)
 				mockClients(controller.cachedClient.(*mockclient.MockClient), controller.azClient, controller.kubeClient)
 				return controller
 			},