Pass the availability zone the Kubernetes node is in when mounting EFS filesystem

pkg/driver/controller.go

@@ -43,6 +43,7 @@ const (
 	DefaultGidMax         = DefaultGidMin + cloud.AccessPointPerFsLimit
 	DefaultTagKey         = "efs.csi.aws.com/cluster"
 	DefaultTagValue       = "true"
+	DiscoverAzName        = "discoverAzName"
 	DirectoryPerms        = "directoryPerms"
 	EnsureUniqueDirectory = "ensureUniqueDirectory"
 	FsId                  = "fileSystemId"
@@ -129,6 +130,7 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 		roleArn                string
 		uid                    int64
 		crossAccountDNSEnabled bool
+		discoverAzNameEnabled  bool
 	)
 
 	//Parse parameters
@@ -143,6 +145,13 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 		return nil, status.Errorf(codes.InvalidArgument, "Missing %v parameter", ProvisioningMode)
 	}
 
+	if value, ok := volumeParams[DiscoverAzName]; ok {
+		discoverAzNameEnabled, err = strconv.ParseBool(value)
+		if err != nil {
+			return nil, status.Errorf(codes.InvalidArgument, "Failed to parse invalid %v: %v", DiscoverAzName, err)
+		}
+	}
+
 	// Create tags
 	tags := map[string]string{
 		DefaultTagKey: DefaultTagValue,
@@ -348,6 +357,10 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 		}
 	}
 
+	if discoverAzNameEnabled {
+		volContext[DiscoverAzName] = strconv.FormatBool(true)
+	}
+
 	return &csi.CreateVolumeResponse{
 		Volume: &csi.Volume{
 			CapacityBytes: volSize,

pkg/driver/node.go

@@ -106,6 +106,15 @@ func (d *Driver) NodePublishVolume(ctx context.Context, req *csi.NodePublishVolu
 			if err != nil {
 				return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("Volume context property %q must be a boolean value: %v", k, err))
 			}
+		case strings.ToLower(DiscoverAzName):
+			discoverAzNameEnabled, err := strconv.ParseBool(v)
+			if err != nil {
+				return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("Volume context property %q must be a boolean value: %v", k, err))
+			}
+			if discoverAzNameEnabled {
+				az := d.cloud.GetMetadata().GetAvailabilityZone()
+				mountOptions = append(mountOptions, "az="+az)
+			}
 		default:
 			return nil, status.Errorf(codes.InvalidArgument, "Volume context property %s not supported.", k)
 		}

pkg/driver/node_test.go

@@ -27,6 +27,7 @@ import (
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/golang/mock/gomock"
+	"github.com/kubernetes-sigs/aws-efs-csi-driver/pkg/cloud"
 	"github.com/kubernetes-sigs/aws-efs-csi-driver/pkg/driver/mocks"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
@@ -42,10 +43,30 @@ type errtyp struct {
 	message string
 }
 
-func setup(mockCtrl *gomock.Controller, volStatter VolStatter, volMetricsOptIn bool) (*mocks.MockMounter, *Driver, context.Context) {
+type cloudMetadata struct {
+	instanceID       string
+	region           string
+	availabilityZone string
+}
+
+var _ cloud.MetadataService = &cloudMetadata{}
+
+func (m *cloudMetadata) GetInstanceID() string {
+	return m.instanceID
+}
+func (m *cloudMetadata) GetRegion() string {
+	return m.region
+}
+func (m *cloudMetadata) GetAvailabilityZone() string {
+	return m.availabilityZone
+}
+
+func setup(mockCtrl *gomock.Controller, volStatter VolStatter, volMetricsOptIn bool) (*mocks.MockCloud, *mocks.MockMounter, *Driver, context.Context) {
+	mockCloud := mocks.NewMockCloud(mockCtrl)
 	mockMounter := mocks.NewMockMounter(mockCtrl)
 	nodeCaps := SetNodeCapOptInFeatures(volMetricsOptIn)
 	driver := &Driver{
+		cloud:           mockCloud,
 		endpoint:        "endpoint",
 		nodeID:          "nodeID",
 		mounter:         mockMounter,
@@ -54,7 +75,7 @@ func setup(mockCtrl *gomock.Controller, volStatter VolStatter, volMetricsOptIn b
 		nodeCaps:        nodeCaps,
 	}
 	ctx := context.Background()
-	return mockMounter, driver, ctx
+	return mockCloud, mockMounter, driver, ctx
 }
 
 func testResult(t *testing.T, funcName string, ret interface{}, err error, expectError errtyp) {
@@ -105,6 +126,7 @@ func TestNodePublishVolume(t *testing.T) {
 		mountSuccess    bool
 		volMetricsOptIn bool
 		expectError     errtyp
+		cloudMetadata   cloud.MetadataService
 	}{
 		{
 			name: "success: normal",
@@ -362,6 +384,21 @@ func TestNodePublishVolume(t *testing.T) {
 			mountArgs:     []interface{}{volumeId + ":/", targetPath, "efs", []string{"tls"}},
 			mountSuccess:  true,
 		},
+		{
+			name: "success: normal with discoverAzName true volume context",
+			req: &csi.NodePublishVolumeRequest{
+				VolumeId:         volumeId,
+				VolumeCapability: stdVolCap,
+				TargetPath:       targetPath,
+				VolumeContext:    map[string]string{"discoverAzName": "true"},
+			},
+			expectMakeDir: true,
+			mountArgs:     []interface{}{volumeId + ":/", targetPath, "efs", []string{"az=us-east-1a", "tls"}},
+			mountSuccess:  true,
+			cloudMetadata: &cloudMetadata{
+				availabilityZone: "us-east-1a",
+			},
+		},
 		{
 			name: "fail: conflicting access point in volume handle and mount options",
 			req: &csi.NodePublishVolumeRequest{
@@ -594,8 +631,11 @@ func TestNodePublishVolume(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			mockCtrl := gomock.NewController(t)
 			defer mockCtrl.Finish()
-			mockMounter, driver, ctx := setup(mockCtrl, NewVolStatter(), tc.volMetricsOptIn)
+			mockCloud, mockMounter, driver, ctx := setup(mockCtrl, NewVolStatter(), tc.volMetricsOptIn)
 
+			if tc.cloudMetadata != nil {
+				mockCloud.EXPECT().GetMetadata().Return(tc.cloudMetadata)
+			}
 			if tc.expectMakeDir {
 				var err error
 				// If not expecting mount, it's because mkdir errored
@@ -722,7 +762,7 @@ func TestNodeUnpublishVolume(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			mockCtrl := gomock.NewController(t)
 			defer mockCtrl.Finish()
-			mockMounter, driver, ctx := setup(mockCtrl, NewVolStatter(), true)
+			_, mockMounter, driver, ctx := setup(mockCtrl, NewVolStatter(), true)
 
 			if tc.expectGetDeviceName {
 				mockMounter.EXPECT().
@@ -849,7 +889,7 @@ func TestNodeGetVolumeStats(t *testing.T) {
 			//setup
 			mockCtrl := gomock.NewController(t)
 			defer mockCtrl.Finish()
-			_, driver, ctx = setup(mockCtrl, NewVolStatter(), true)
+			_, _, driver, ctx = setup(mockCtrl, NewVolStatter(), true)
 
 			if tc.updateCache {
 				mu.Lock()